import asyncio
import logging
from pathlib import Path

from sqlalchemy.ext.asyncio import async_sessionmaker

from kbqa.api.errors import AppError
from kbqa.config import Settings
from kbqa.documents.models import Chunk
from kbqa.documents.repository import DocumentRepository
from kbqa.indexing.loaders import load_document
from kbqa.indexing.splitter import split_pages
from kbqa.rag.embeddings import EmbeddingProvider
from kbqa.rag.types import PreparedChunk
from kbqa.rag.vector_store import MilvusVectorStore

logger = logging.getLogger(__name__)


class IndexWorker:
    def __init__(
        self,
        settings: Settings,
        session_factory: async_sessionmaker,
        embeddings: EmbeddingProvider,
        vector_store: MilvusVectorStore,
        vector_mutation_lock: asyncio.Lock,
    ) -> None:
        self.settings = settings
        self.session_factory = session_factory
        self.embeddings = embeddings
        self.vector_store = vector_store
        self.vector_mutation_lock = vector_mutation_lock
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self._task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        async with self.session_factory() as session:
            pending_ids = await DocumentRepository(session).requeue_unfinished()
        for job_id in pending_ids:
            self.queue.put_nowait(job_id)
        self._task = asyncio.create_task(self._run(), name="kbqa-index-worker")

    async def stop(self) -> None:
        if self._task is None:
            return
        self._task.cancel()
        try:
            await self._task
        except asyncio.CancelledError:
            pass
        self._task = None

    def enqueue(self, job_id: str) -> None:
        self.queue.put_nowait(job_id)

    async def _run(self) -> None:
        while True:
            job_id = await self.queue.get()
            try:
                await self._process(job_id)
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Unexpected indexing worker failure job_id=%s", job_id)
            finally:
                self.queue.task_done()

    async def _process(self, job_id: str) -> None:
        async with self.session_factory() as session:
            repository = DocumentRepository(session)
            job = await repository.claim_job(job_id)
            if job is None:
                return
            document = await repository.get(job.document_id)
            if document is None:
                return
            document_id = document.id
            stored_path = document.stored_path
            file_type = document.file_type
        try:
            async with self.vector_mutation_lock:
                await asyncio.to_thread(self.vector_store.delete_document, document_id)
            async with self.session_factory() as session:
                await DocumentRepository(session).replace_chunks(document_id, [])
            pages = await asyncio.to_thread(load_document, Path(stored_path), file_type)
            prepared = split_pages(
                document_id,
                pages,
                self.settings.chunk_size,
                self.settings.chunk_overlap,
            )
            if not prepared:
                raise AppError("EMPTY_DOCUMENT", "文档分块结果为空", 415)
            orm_chunks = [
                Chunk(
                    id=chunk.id,
                    document_id=chunk.document_id,
                    order_index=chunk.order_index,
                    content=chunk.content,
                    page_number=chunk.page_number,
                    char_count=len(chunk.content),
                )
                for chunk in prepared
            ]
            async with self.session_factory() as session:
                await DocumentRepository(session).replace_chunks(document_id, orm_chunks)
            completed = await self._embed_and_upsert(document_id, prepared)
            if not completed:
                return
            async with self.session_factory() as session:
                await DocumentRepository(session).mark_indexed(document_id, len(prepared))
        except asyncio.CancelledError:
            raise
        except AppError as exc:
            await self._mark_failed(document_id, exc.code, exc.message)
        except Exception:
            logger.exception("Document indexing failed document_id=%s", document_id)
            await self._mark_failed(document_id, "INDEXING_FAILED", "文档索引失败，请重试")

    async def _embed_and_upsert(self, document_id: str, chunks: list[PreparedChunk]) -> bool:
        batch_size = 10
        for start in range(0, len(chunks), batch_size):
            batch = chunks[start : start + batch_size]
            vectors = await self.embeddings.embed_documents([chunk.content for chunk in batch])
            async with self.vector_mutation_lock:
                async with self.session_factory() as session:
                    if not await DocumentRepository(session).is_indexing(document_id):
                        return False
                await asyncio.to_thread(self.vector_store.upsert, batch, vectors)
        return True

    async def _mark_failed(self, document_id: str, code: str, message: str) -> None:
        async with self.session_factory() as session:
            await DocumentRepository(session).mark_failed(document_id, code, message)