from dataclasses import dataclass

from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from kbqa.database import utc_now
from kbqa.documents.models import Chunk, Document, IndexJob, new_uuid


@dataclass(frozen=True, slots=True)
class StoredUpload:
    document_id: str
    filename: str
    stored_path: str
    file_type: str
    mime_type: str
    file_size: int
    sha256: str


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

    async def create_pending(self, upload: StoredUpload) -> tuple[Document, IndexJob]:
        document = Document(
            id=upload.document_id,
            filename=upload.filename,
            stored_path=upload.stored_path,
            file_type=upload.file_type,
            mime_type=upload.mime_type,
            file_size=upload.file_size,
            sha256=upload.sha256,
            status="pending",
        )
        job = IndexJob(id=new_uuid(), document_id=document.id, status="queued")
        self.session.add_all([document, job])
        await self.session.commit()
        await self.session.refresh(document)
        await self.session.refresh(job)
        return document, job

    async def page(
        self,
        status: str | None,
        offset: int,
        limit: int,
    ) -> tuple[list[Document], int]:
        filters = [Document.status == status] if status else []
        total = await self.session.scalar(select(func.count(Document.id)).where(*filters))
        result = await self.session.scalars(
            select(Document)
            .where(*filters)
            .order_by(Document.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        return list(result), int(total or 0)

    async def get(self, document_id: str) -> Document | None:
        return await self.session.get(Document, document_id)

    async def is_indexing(self, document_id: str) -> bool:
        status = await self.session.scalar(
            select(Document.status).where(Document.id == document_id)
        )
        return status == "indexing"

    async def get_job_for_document(self, document_id: str) -> IndexJob | None:
        return await self.session.scalar(
            select(IndexJob).where(IndexJob.document_id == document_id)
        )

    async def page_chunks(
        self,
        document_id: str,
        offset: int,
        limit: int,
    ) -> tuple[list[Chunk], int]:
        total = await self.session.scalar(
            select(func.count(Chunk.id)).where(Chunk.document_id == document_id)
        )
        result = await self.session.scalars(
            select(Chunk)
            .where(Chunk.document_id == document_id)
            .order_by(Chunk.order_index)
            .offset(offset)
            .limit(limit)
        )
        return list(result), int(total or 0)

    async def claim_job(self, job_id: str) -> IndexJob | None:
        job = await self.session.get(IndexJob, job_id)
        if job is None or job.status != "queued":
            return None
        document = await self.session.get(Document, job.document_id)
        if document is None or document.status not in {"pending", "failed"}:
            return None
        now = utc_now()
        job.status = "running"
        job.attempt_count += 1
        job.started_at = now
        job.finished_at = None
        job.last_error = None
        document.status = "indexing"
        document.error_code = None
        document.error_message = None
        document.updated_at = now
        await self.session.commit()
        return job

    async def replace_chunks(self, document_id: str, chunks: list[Chunk]) -> None:
        await self.session.execute(delete(Chunk).where(Chunk.document_id == document_id))
        self.session.add_all(chunks)
        await self.session.commit()

    async def mark_indexed(self, document_id: str, chunk_count: int) -> None:
        document = await self.session.get(Document, document_id)
        job = await self.get_job_for_document(document_id)
        if document is None or job is None:
            return
        now = utc_now()
        document.status = "indexed"
        document.chunk_count = chunk_count
        document.updated_at = now
        document.error_code = None
        document.error_message = None
        job.status = "succeeded"
        job.finished_at = now
        job.last_error = None
        await self.session.commit()

    async def mark_failed(self, document_id: str, code: str, message: str) -> None:
        document = await self.session.get(Document, document_id)
        job = await self.get_job_for_document(document_id)
        if document is None or job is None:
            return
        now = utc_now()
        document.status = "failed"
        document.error_code = code
        document.error_message = message
        document.updated_at = now
        job.status = "failed"
        job.finished_at = now
        job.last_error = message
        await self.session.commit()

    async def requeue_unfinished(self) -> list[str]:
        jobs = list(
            await self.session.scalars(
                select(IndexJob).where(IndexJob.status.in_(("queued", "running")))
            )
        )
        now = utc_now()
        for job in jobs:
            job.status = "queued"
            job.started_at = None
            job.finished_at = None
            document = await self.session.get(Document, job.document_id)
            if document is not None:
                document.status = "pending"
                document.updated_at = now
        await self.session.commit()
        return [job.id for job in jobs]

    async def requeue_failed(self, document_id: str) -> IndexJob | None:
        document = await self.session.get(Document, document_id)
        job = await self.get_job_for_document(document_id)
        if document is None or job is None or document.status != "failed":
            return None
        document.status = "pending"
        document.error_code = None
        document.error_message = None
        document.updated_at = utc_now()
        job.status = "queued"
        job.last_error = None
        job.finished_at = None
        await self.session.commit()
        return job

    async def mark_deleting(self, document_id: str) -> Document | None:
        document = await self.session.get(Document, document_id)
        if document is None:
            return None
        document.status = "deleting"
        document.updated_at = utc_now()
        await self.session.commit()
        return document

    async def deleting_documents(self) -> list[Document]:
        return list(
            await self.session.scalars(select(Document).where(Document.status == "deleting"))
        )

    async def delete_document(self, document_id: str) -> bool:
        document = await self.session.get(Document, document_id)
        if document is None:
            return False
        await self.session.delete(document)
        await self.session.commit()
        return True
