"""Orchestrates: understand question -> retrieve -> ask the LLM -> return answer + sources."""
import logging

from . import llm, query, retrieval
from .config import settings
from .registry import get_registry

log = logging.getLogger("rag.rag")

SYSTEM_PROMPT = """You answer questions about the user's uploaded documents using ONLY the provided context.
- Cite sources inline as [1], [2] using the numbers in the context.
- If the context does not contain the answer, say clearly that you could not find it in the documents. Never guess.
- If the user asks what is on a specific page, describe or summarise that page from the context.
- If the context lists notes about missing pages, mention them briefly.
- Reply in the same language the user wrote in (English, Hindi or Hinglish).
- Be concise and well organised."""


class NoReadyDocuments(Exception):
    pass


def _merge_pages(chunks: list[dict]) -> list[dict]:
    """For page lookups: one source per (document, page) instead of one per chunk."""
    merged: dict[tuple, dict] = {}
    for c in chunks:
        key = (c["doc_id"], c["page"])
        if key in merged:
            merged[key]["text"] += "\n\n" + c["text"]
        else:
            merged[key] = {**c}
    return list(merged.values())


def _page_notes(scope: list[dict], pages: list[int], found: set[tuple]) -> list[str]:
    notes = []
    for d in scope:
        for p in pages:
            if (d["doc_id"], p) in found:
                continue
            if d.get("pages") and p > d["pages"]:
                notes.append(f'"{d["filename"]}" has only {d["pages"]} pages, so page {p} does not exist.')
            else:
                notes.append(
                    f'Page {p} of "{d["filename"]}" has no extractable text (blank page or a scanned image).'
                )
    return notes


def _build_context(sources: list[dict]) -> str:
    parts, used = [], 0
    for i, s in enumerate(sources):
        block = f"[{i + 1}] ({s['filename']}, page {s['page']})\n{s['text']}"
        room = settings.max_context_chars - used
        if room <= 200:
            break
        block = block[:room]
        parts.append(block)
        used += len(block)
    return "\n\n".join(parts)


def _shape(sources: list[dict]) -> list[dict]:
    return [
        {
            "doc_id": s["doc_id"],
            "filename": s["filename"],
            "page": s["page"],
            "score": round(float(s.get("score", 0.0)), 4),
            "text": s["text"],
        }
        for s in sources
    ]


def answer_question(question: str, doc_ids: list[str] | None = None, history: list[dict] | None = None) -> dict:
    history = history or []
    ready = get_registry().ready()
    if not ready:
        raise NoReadyDocuments()

    scope = [d for d in ready if d["doc_id"] in doc_ids] if doc_ids else ready
    if not scope:
        scope = ready

    plan = query.analyze(question, scope, history)
    explicit = len(scope) == 1
    if plan.doc_ids:  # the question names a document -> narrow to it
        scope = [d for d in scope if d["doc_id"] in plan.doc_ids] or scope
        explicit = True
    scope_ids = [d["doc_id"] for d in scope]

    notes: list[str] = []
    note_for_user = ""
    if plan.intent == "page_lookup" and plan.pages:
        pages = plan.pages
        if len(pages) > settings.max_lookup_pages:
            note_for_user = f"\n\n(I read the first {settings.max_lookup_pages} of the {len(pages)} pages you asked about.)"
            pages = pages[: settings.max_lookup_pages]
        chunks = retrieval.fetch_pages(scope_ids, pages)
        sources = _merge_pages(chunks)
        found = {(s["doc_id"], s["page"]) for s in sources}
        notes = _page_notes(scope, pages, found)
        if sources and len(scope) > 1 and not explicit:
            notes = []  # don't nag about documents that simply don't have that page
        mode = "page_lookup"
        shown_pages = pages
    else:
        sources = retrieval.hybrid_search(plan.search_query, scope_ids, plan.pages or None)
        mode = "search"
        shown_pages = plan.pages

    meta = {
        "mode": mode,
        "search_query": plan.search_query if mode == "search" else None,
        "pages": shown_pages or None,
        "documents": [d["filename"] for d in scope],
    }

    if not sources:
        msg = " ".join(notes) if notes else "I couldn't find anything relevant to this in the selected documents."
        return {"answer": msg, "sources": [], **meta}

    context = _build_context(sources)
    note_block = ("\n\nNotes:\n- " + "\n- ".join(notes)) if notes else ""
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history[-4:]:
        messages.append({"role": h["role"], "content": h["content"][:600]})
    messages.append({"role": "user", "content": f"Context:\n{context}{note_block}\n\nQuestion: {question}"})

    answer = llm.chat(messages)
    if not answer:
        answer = "The model returned an empty reply. Please try asking again."
    return {"answer": answer + note_for_user, "sources": _shape(sources), **meta}