from sqlalchemy import create_engine, select, func, Column, Integer, String, Text, Float, BigInteger
from sqlalchemy.orm import sessionmaker, Session, declarative_base
from sqlalchemy.pool import QueuePool
import sqlalchemy
from datetime import datetime
from tqdm import tqdm
from enum import Enum
from faker import Faker
import random
import os
from contextlib import contextmanager
from typing import List, Optional, Any
from dataclasses import dataclass

engine = create_engine(
    # 'sqlite:///rest_server_books.db',
    'postgresql://postgres:liuyanfeng66@localhost:5432/BookRestGrpcCompare',
    poolclass=QueuePool,
    pool_size=30,
    max_overflow=50,
    pool_pre_ping=True,
    echo=False,
    # connect_args={
    #     "check_same_thread": False,
    #     "timeout": 30  # 文件锁最长等待秒数，视情况可调
    # }
)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()


@dataclass
class PaginationResult:
    """分页结果数据类"""
    items: List[Any]
    total: int
    page: int
    per_page: int
    pages: int
    has_prev: bool
    has_next: bool
    prev_num: Optional[int]
    next_num: Optional[int]


class Paginator:
    def __init__(self, session: Session):
        self.session = session

    def paginate(self, query, page: int = 1, per_page: int = 20, error_out: bool = True) -> PaginationResult:
        if per_page <= 0:
            return PaginationResult(
                items=[],
                total=0,
                page=1,
                per_page=0,
                pages=0,
                has_prev=False,
                has_next=False,
                prev_num=None,
                next_num=None
            )

        page = max(1, page)
        per_page = max(1, per_page)

        # Get total count of items
        count_query = select(func.count()).select_from(query.statement.alias())
        total = self.session.scalar(count_query)

        # Get total pages
        pages = (total + per_page - 1) // per_page

        if error_out and page > pages and total > 0:
            raise ValueError(f"Page {page} out of range (1-{pages})")

        offset = (page - 1) * per_page

        items = query.offset(offset).limit(per_page).all()

        return PaginationResult(
            items=items,
            total=total,
            page=page,
            per_page=per_page,
            pages=pages,
            has_prev=page > 1,
            has_next=page < pages,
            prev_num=page - 1 if page > 1 else None,
            next_num=page + 1 if page < pages else None
        )


class BaseModel(Base):
    __abstract__ = True

    def to_dict(self, exclude: set = None):
        exclude = exclude or set()
        ret = {}
        for c in self.__table__.columns:
            if c.name not in exclude:
                curr_item = getattr(self, c.name)
                if isinstance(curr_item, Enum):
                    ret[c.name] = curr_item.value
                else:
                    ret[c.name] = curr_item

        return ret


class Books(BaseModel):
    __tablename__ = 'books'
    id = Column(BigInteger, primary_key=True, autoincrement=True)
    isbn: str = Column(String(20), nullable=True, index=True)
    barcode: str = Column(String(50), nullable=True)

    title: str = Column(Text, nullable=False)
    subtitle: str = Column(Text, nullable=True)
    author: str = Column(Text, nullable=False)
    translator: str = Column(Text, nullable=True)
    editor: str = Column(Text, nullable=True)
    publisher: str = Column(String(200), nullable=True)
    publication_date = Column(BigInteger, nullable=True)
    edition: str = Column(String(50), nullable=True)
    pages: int = Column(Integer, nullable=True)
    language: str = Column(String(50), nullable=True)

    category_id: int = Column(Integer, nullable=True, index=True)
    subject: str = Column(String(200), nullable=True)
    keywords: str = Column(Text, nullable=True)
    description: str = Column(Text, nullable=True)
    abstract: str = Column(Text, nullable=True)

    format: str = Column(String(50), nullable=True)
    binding: str = Column(String(20), nullable=True)
    weight: float = Column(Float, nullable=True)
    cover_image: str = Column(String(500), nullable=True)


def generate_random_book_data(book_id: int, fake: Faker) -> dict:
    return {
        'id': book_id,
        'isbn': fake.isbn13(separator='-'),
        'barcode': fake.ean(length=13),
        'title': fake.sentence(nb_words=6),
        'subtitle': fake.sentence(nb_words=3),
        'author': fake.name(),
        'translator': fake.name(),
        'editor': fake.name(),
        'publisher': fake.company(),
        'publication_date': random.randint(-2190472676, 1754206067),
        'edition': fake.word(),
        'pages': random.randint(100, 1000),
        'language': random.choice(['English', 'Français', 'Deutsch']),
        'category_id': random.randint(1, 100),
        'subject': fake.word(),
        'keywords': ", ".join(fake.words(nb=5, unique=True, ext_word_list=None)),
        'description': fake.text(max_nb_chars=200),
        'abstract': fake.text(max_nb_chars=100),
        'format': random.choice(['16mo', '8vo', '4to', 'Folio', 'Quarto']),
        'binding': random.choice(['expensive', 'cheap', 'paperback', 'hardcover']),
        'weight': random.uniform(100, 2000),
        'cover_image': "https://example.com/cover_image.jpg"
    }


@contextmanager
def get_db_session():
    session = SessionLocal()
    try:
        yield session
    except Exception as e:
        session.rollback()
        print(f"Database error: {e}")
        raise
    finally:
        session.close()


def check_and_generate_random_data(generate_count=10000):
    Base.metadata.create_all(bind=engine)

    fake = Faker('en_GB')
    batch_size = 10000

    with get_db_session() as session:
        data_count = session.query(Books).count()
        if data_count != 0:
            if input(
                    f"Database already contains data ({data_count}). Do you want to regenerate random data? (y/n): ").lower() != 'y':
                return

            if input("Delete all existing data? (y/n): ").lower() == 'y':
                if data_count > 0:
                    session.query(Books).delete()
                    session.commit()

    for batch_start in tqdm(range(0, generate_count, batch_size), desc="Inserting batches"):
        batch_end = min(batch_start + batch_size, generate_count)

        batch_data = [
            generate_random_book_data(i, fake)
            for i in range(batch_start, batch_end)
        ]

        with get_db_session() as session:
            session.bulk_insert_mappings(Books, batch_data)
            session.commit()

#
# with app.app_context():
#     check_and_generate_random_data()

# check_and_generate_random_data()


if __name__ == "__main__":
    check_and_generate_random_data(int(input("Enter the number of random books to generate: ")))
    print("Database initialized and random data generated.")
