"""Ask pipeline: question -> MaxSim retrieval over the store -> top-K page
images -> MiniCPM answer grounded in those pages."""

from __future__ import annotations

from core.pdf import render_page
from core.store import Store
from models.colembed import ColEmbed
from models.minicpm import MiniCPM


class AskPipeline:
    def __init__(self, embedder: ColEmbed, store: Store, llm: MiniCPM):
        self.embedder = embedder
        self.store = store
        self.llm = llm

    def run(self, question: str, doc_ids: list[str] | None, top_k: int):
        """Return (answer markdown, gallery items [(image, caption)])."""
        question = (question or "").strip()
        if not question:
            raise ValueError("Please enter a question.")
        docs = self.store.list_docs()
        if not docs:
            raise ValueError("No manuals indexed yet — add one in the Library tab.")

        hits = self.embedder.search(question, self.store, doc_ids or None, int(top_k))
        names = {d["doc_id"]: d["name"] for d in docs}
        pages = [
            (f"{names[doc_id]} — p.{page}", render_page(self.store.pdf_path(doc_id), page), score)
            for doc_id, page, score in hits
        ]
        answer = self.llm.answer(question, [(label, img) for label, img, _ in pages])
        gallery = [(img, f"{label} (score {score:.1f})") for label, img, score in pages]
        return answer, gallery