import asyncio

from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker

from kbqa.documents.models import Chunk, Document
from kbqa.rag.embeddings import EmbeddingProvider
from kbqa.rag.types import SourceCandidate
from kbqa.rag.vector_store import MilvusVectorStore


class KnowledgeRetriever:
    def __init__(
        self,
        embeddings: EmbeddingProvider,
        vector_store: MilvusVectorStore,
        session_factory: async_sessionmaker,
    ) -> None:
        self.embeddings = embeddings
        self.vector_store = vector_store
        self.session_factory = session_factory

    async def retrieve(self, query: str, top_k: int) -> list[SourceCandidate]:
        vector = await self.embeddings.embed_query(query)
        candidate_limit = max(top_k * 3, top_k)
        hits = await asyncio.to_thread(self.vector_store.search, vector, candidate_limit)
        if not hits:
            return []
        ids = [hit.chunk_id for hit in hits]
        async with self.session_factory() as session:
            rows = await session.execute(
                select(Chunk, Document)
                .join(Document, Document.id == Chunk.document_id)
                .where(Chunk.id.in_(ids), Document.status == "indexed")
            )
            lookup = {chunk.id: (chunk, document) for chunk, document in rows}
        sources: list[SourceCandidate] = []
        for hit in hits:
            record = lookup.get(hit.chunk_id)
            if record is None:
                continue
            chunk, document = record
            sources.append(
                SourceCandidate(
                    chunk_id=chunk.id,
                    document_id=document.id,
                    filename=document.filename,
                    order_index=chunk.order_index,
                    page_number=chunk.page_number,
                    content=chunk.content,
                    score=hit.score,
                )
            )
            if len(sources) == top_k:
                break
        return sources