"""Retrieval with citations (build plan Phase 3).

Embeds the query and searches the workspace's knowledge-base collections,
returning ranked passages with their source document for citation.
"""

from __future__ import annotations

import uuid
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.orm import Session

from legate.db.models.knowledge_base import KnowledgeBase
from legate.rag.embeddings.base import Embedder
from legate.rag.vectorstore.base import VectorStore


@dataclass
class Citation:
    document_id: str
    knowledge_base_id: str
    source: str
    text: str
    score: float
    chunk_idx: int

    def as_dict(self) -> dict:
        return {
            "document_id": self.document_id,
            "knowledge_base_id": self.knowledge_base_id,
            "source": self.source,
            "text": self.text,
            "score": round(self.score, 6),
            "chunk_idx": self.chunk_idx,
        }


def _knowledge_bases(
    db: Session, workspace_id: uuid.UUID, knowledge_base_id: uuid.UUID | None
) -> list[KnowledgeBase]:
    stmt = select(KnowledgeBase).where(KnowledgeBase.workspace_id == workspace_id)
    if knowledge_base_id is not None:
        stmt = stmt.where(KnowledgeBase.id == knowledge_base_id)
    return list(db.scalars(stmt).all())


def search(
    *,
    db: Session,
    workspace_id: uuid.UUID,
    query: str,
    embedder: Embedder,
    vector_store: VectorStore,
    top_k: int = 5,
    knowledge_base_id: uuid.UUID | None = None,
) -> list[Citation]:
    query = query.strip()
    if not query:
        return []
    kbs = _knowledge_bases(db, workspace_id, knowledge_base_id)
    if not kbs:
        return []

    query_vector = embedder.embed([query])[0]
    citations: list[Citation] = []
    for kb in kbs:
        for hit in vector_store.search(kb.vector_collection, query_vector, top_k=top_k):
            payload = hit.payload or {}
            citations.append(
                Citation(
                    document_id=str(payload.get("document_id", "")),
                    knowledge_base_id=str(kb.id),
                    source=str(payload.get("source", "")),
                    text=str(payload.get("text", "")),
                    score=hit.score,
                    chunk_idx=int(payload.get("chunk_idx", 0)),
                )
            )
    citations.sort(key=lambda c: c.score, reverse=True)
    return citations[:top_k]
