"""Document 文档 Repository。"""

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.models.document import Document


class DocumentRepository:
    def __init__(self, session: Session) -> None:
        self._session = session

    def get_by_id(self, doc_id: str) -> Document | None:
        return self._session.get(Document, doc_id)

    def get_by_doc_token_hash(self, token_hash: str) -> Document | None:
        stmt = select(Document).where(Document.doc_token_hash == token_hash)
        return self._session.scalars(stmt).first()

    def list_by_knowledge_base(
        self,
        kb_id: str,
        *,
        category_id: str | None = None,
        page: int = 1,
        page_size: int = 50,
        status: str | None = None,
    ) -> tuple[list[Document], int]:
        """分页获取知识库的文档列表。支持按分类过滤。"""
        conditions = [
            Document.knowledge_base_id == kb_id,
            Document.status != "DELETED",
        ]
        if status:
            conditions.append(Document.status == status)
        if category_id:
            conditions.append(Document.category_id == category_id)

        count_stmt = select(func.count()).select_from(Document).where(*conditions)
        total = self._session.scalar(count_stmt) or 0

        stmt = (
            select(Document)
            .where(*conditions)
            .order_by(Document.created_at.desc())
            .offset((page - 1) * page_size)
            .limit(page_size)
        )
        items = list(self._session.scalars(stmt).all())
        return items, total

    def list_by_user(
        self, user_id: str, *, page: int = 1, page_size: int = 20
    ) -> tuple[list[Document], int]:
        """分页获取用户的文档列表。"""
        conditions = [Document.user_id == user_id, Document.status != "DELETED"]

        count_stmt = select(func.count()).select_from(Document).where(*conditions)
        total = self._session.scalar(count_stmt) or 0

        stmt = (
            select(Document)
            .where(*conditions)
            .order_by(Document.created_at.desc())
            .offset((page - 1) * page_size)
            .limit(page_size)
        )
        items = list(self._session.scalars(stmt).all())
        return items, total

    def create(
        self,
        *,
        knowledge_base_id: str,
        user_id: str,
        original_filename: str,
        storage_path: str,
        file_size: int,
        mime_type: str,
        file_ext: str,
        sha256: str,
        doc_token_hash: str,
        doc_token_encrypted: str,
        doc_token_hint: str,
        category_id: str | None = None,
    ) -> Document:
        doc = Document(
            knowledge_base_id=knowledge_base_id,
            user_id=user_id,
            original_filename=original_filename,
            storage_path=storage_path,
            file_size=file_size,
            mime_type=mime_type,
            file_ext=file_ext,
            sha256=sha256,
            doc_token_hash=doc_token_hash,
            doc_token_encrypted=doc_token_encrypted,
            doc_token_hint=doc_token_hint,
            category_id=category_id,
            status="PENDING",
        )
        self._session.add(doc)
        self._session.flush()
        return doc

    def update(self, doc: Document, **fields) -> None:
        for key, value in fields.items():
            setattr(doc, key, value)
        self._session.flush()

    def delete(self, doc: Document) -> None:
        doc.status = "DELETED"
        self._session.flush()

    def count_by_user(self, user_id: str) -> int:
        stmt = (
            select(func.count())
            .select_from(Document)
            .where(Document.user_id == user_id, Document.status != "DELETED")
        )
        return self._session.scalar(stmt) or 0