from __future__ import annotations

import json
import re
from dataclasses import dataclass, replace

from . import grounding, metrics as metric_names
from .metrics import Metrics, NullMetrics
from .models import EvidenceDecision, QueryIntent, RetrievalResult, SubjectScope
from .ports import AnswerGenerationUnavailable, AnswerGenerator
from .prompt import build_entailment_request, build_request, build_sufficiency_request
from .routing import QueryRoutingService


@dataclass(frozen=True)
class Citation:
    chunk_id: str
    printed_page_start: int
    printed_page_end: int
    physical_page: int
    block_id: str | None = None
    bbox: tuple[float, float, float, float] | None = None
    source_crop: str | None = None
    attachment: str | None = None


@dataclass(frozen=True)
class GroundedAnswer:
    result: RetrievalResult
    answer: str | None
    citations: tuple[Citation, ...] = ()
    generated: bool = False
    # Set when the model decided the turn is under-specified and asked back
    # (e.g. a dose question with no age/weight). The answer field carries the
    # question; the caller renders it as a clarification, not a final answer.
    clarification: str | None = None


@dataclass(frozen=True)
class _GenOutcome:
    answer: str | None = None
    clarification: str | None = None


class GroundedAnswerService:
    """Retrieval decides what is true; generation only decides how it reads.

    Two operating modes, not to be confused with each other:

    - **No generator configured** (`generator=None`, e.g. `ANSWER_PROVIDER=
      disabled`) is retrieval-only mode, a deliberate and fully supported
      way to run this service. It quotes the retrieved source verbatim.
    - **A generator IS configured.** Its output replaces the extractive text
      only if it clears two independent checks: `grounding.verify` (every
      figure and citation traces to the specific evidence block it cites,
      and every claim carries one) and `_verify_entailment` (a second LLM
      pass confirming each cited claim's *content* — not just its numbers —
      is actually stated by that block). If a configured generation fails
      any check, or the provider itself is unreachable, or its output is
      malformed, the turn **abstains** (`reason="generation_unavailable"`
      or the specific `grounding.verify` reason) rather than silently
      degrading to a raw source dump — this product is a real LLM chatbot,
      and a citation-stapled paragraph of book text is not an acceptable
      stand-in for an answer the model was supposed to produce.
    """

    def __init__(
        self,
        routing: QueryRoutingService,
        generator: AnswerGenerator | None = None,
        metrics: Metrics | None = None,
    ) -> None:
        self._routing = routing
        self._generator = generator
        self._metrics = metrics or NullMetrics()

    def answer(
        self,
        query: str,
        subject_scope: SubjectScope,
        intent: QueryIntent,
        drug_id: str | None = None,
    ) -> GroundedAnswer:
        # When the caller already resolved the drug (e.g. the conversational
        # layer, incl. an inherited follow-up), retrieve for it directly instead
        # of re-resolving from the turn text — re-resolution from a rewritten
        # turn is what abstained good follow-ups as "ambiguous".
        if drug_id is not None:
            result = self._routing.retrieve_for_drug(
                query, drug_id, subject_scope, intent
            )
        else:
            result = self._routing.retrieve(query, subject_scope, intent)
        return self.answer_from_result(query, result)

    def answer_from_result(
        self, query: str, result: RetrievalResult
    ) -> GroundedAnswer:
        """Everything after retrieval — grounding, sufficiency, generation,
        citations. Split out so the new understanding-driven orchestrator
        (`rag/agent.py`) reuses the safe answer path without going through the
        old `QueryRoutingService` text resolution."""
        if result.decision == EvidenceDecision.ABSTAIN:
            self._metrics.increment(metric_names.ABSTENTION, reason=result.reason)
            return GroundedAnswer(result, None)

        indexed = self._indexed_citations(result)
        if indexed is None:
            return GroundedAnswer(
                replace(
                    result,
                    decision=EvidenceDecision.ABSTAIN,
                    reason="missing_printed_page_provenance",
                    evidence=(),
                ),
                None,
            )
        all_citations = tuple(citation for _, citation in indexed)
        if result.decision == EvidenceDecision.VERIFY_PDF:
            # Never generated over. A quarantined table or formula is exactly
            # the evidence whose numbers were not reliably reconstructed, so
            # rephrasing it is the one case where fluency could invent a dose.
            return GroundedAnswer(
                result,
                "Nguồn có bảng hoặc công thức cần đối chiếu trực tiếp với ảnh PDF; "
                "không tự động trích số liệu.",
                all_citations,
            )

        evidence_texts = tuple(item.text for item in result.evidence)
        extractive = "\n\n".join(
            f"{text} [{index}]" for index, text in enumerate(evidence_texts, start=1)
        )

        # Reasoning step BEFORE answering: if the turn is under-specified (a dose
        # with several bands and no age/weight/condition), ask instead of dumping.
        # A separate focused call is more reliable than folding it into generation.
        clarify_q = self._check_sufficiency(query, evidence_texts, result.is_drug_overview)
        if clarify_q is not None:
            return GroundedAnswer(result, clarify_q, (), clarification=clarify_q)

        outcome = self._generate(query, evidence_texts, intro=result.is_drug_overview)
        if outcome.clarification is not None:
            # The model judged the turn under-specified (a dose with no
            # age/weight/renal-function/indication…) and asked back instead of
            # listing every band. Return the question, not the whole section.
            return GroundedAnswer(
                result, outcome.clarification, (), clarification=outcome.clarification
            )

        if outcome.answer is None:
            if self._generator is None:
                # No generator configured at all — retrieval-only mode. A
                # deliberate operating mode (e.g. ANSWER_PROVIDER=disabled),
                # not a failure, so the source is quoted verbatim.
                self._metrics.increment(metric_names.ANSWER_EXTRACTIVE)
                citations = self._cited_only(indexed, extractive) or all_citations
                return GroundedAnswer(result, extractive, citations)
            # A generator WAS configured and this specific generation did not
            # clear the safety checks (provider outage, malformed output, an
            # ungrounded/uncited/unsupported claim). This product is a real
            # LLM chatbot, not the retired offline-extractive build — a raw
            # source dump is not an acceptable stand-in for a failed
            # generation, so this abstains instead of silently degrading to
            # one.
            self._metrics.increment(
                metric_names.ABSTENTION, reason="generation_unavailable"
            )
            return GroundedAnswer(
                replace(
                    result,
                    decision=EvidenceDecision.ABSTAIN,
                    reason="generation_unavailable",
                ),
                None,
            )

        # Show only the sources the answer actually cited, not every chunk that
        # was retrieved — a paragraph that cites [4] must not drag 13 citation
        # chips onto the screen. Falls back to all when the text cites nothing.
        citations = self._cited_only(indexed, outcome.answer) or all_citations
        self._metrics.increment(metric_names.GENERATION_SERVED)
        return GroundedAnswer(result, outcome.answer, citations, generated=True)

    def _generate(
        self, query: str, evidence_texts: tuple[str, ...], intro: bool = False
    ) -> "_GenOutcome":
        """A verified generation, a clarifying question, or empty to fall back."""
        if self._generator is None or not evidence_texts:
            return _GenOutcome()

        request = build_request(query, evidence_texts, intro=intro)
        try:
            raw = self._generator.generate(request.system, request.user, request.schema)
        except AnswerGenerationUnavailable:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="provider_unavailable"
            )
            return _GenOutcome()

        try:
            payload = json.loads(raw)
            answer = payload["answer"]
            sufficient = payload["evidence_sufficient"]
        except (ValueError, TypeError, KeyError):
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="malformed_output"
            )
            return _GenOutcome()

        # The model asked for a missing detail (age/weight/renal function/
        # indication…) instead of listing every band. A clarify is not a grounded
        # claim, so it skips the number check — it states no dose.
        clarify = payload.get("clarifying_question") if isinstance(payload, dict) else None
        if isinstance(clarify, str) and clarify.strip():
            return _GenOutcome(clarification=clarify.strip())

        if not isinstance(answer, str) or not isinstance(sufficient, bool):
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="malformed_output"
            )
            return _GenOutcome()
        if not sufficient:
            # The model says the evidence does not answer the question. Showing
            # the retrieved section verbatim lets the clinician judge that.
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="evidence_insufficient"
            )
            return _GenOutcome()

        report = grounding.verify(answer, evidence_texts)
        if not report.grounded:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason=report.reason
            )
            return _GenOutcome()

        if not self._verify_entailment(answer, evidence_texts):
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="unsupported_claim"
            )
            return _GenOutcome()
        return _GenOutcome(answer=answer)

    def _verify_entailment(self, answer: str, evidence_texts: tuple[str, ...]) -> bool:
        """A second, adversarial LLM pass over an answer that already passed
        `grounding.verify`.

        The regex check above only binds numbers and citation indices — it
        has no notion of meaning, so "Metformin chữa ung thư [1]" citing an
        evidence block about "điều trị đái tháo đường" sails through it
        untouched: right drug, right citation shape, fabricated indication.
        This call is what catches that: each substantive, validly-cited claim
        is checked against only the evidence block(s) it names, by a model
        told to compare wording, not to reason about medicine.

        Fails closed on an outage or malformed output. A single rejection is
        NOT: live probing (2026-08-06) found the judge call itself is noisy
        — the identical claim/evidence pair, called three times, came back
        entailed twice and rejected once, discarding a correct, well-cited
        interaction answer. So a reject triggers one same-claim retry, and
        only a second, agreeing reject discards the generation; a single
        provider outage/malformed response still fails closed immediately
        (that failure mode is reliable, not noisy — no retry needed there).
        An answer with no claim text at all (nothing between or after its
        citation markers) is vacuously fine — nothing to verify, no call.
        """
        claims = [
            (claim.text, "\n".join(evidence_texts[i - 1] for i in claim.indices))
            for claim in grounding.split_claims(answer, len(evidence_texts))
            if claim.indices and grounding.has_content(claim.text)
        ]
        if not claims:
            return True

        request = build_entailment_request(claims)
        first = self._run_entailment_check(request)
        if first is None:
            return False
        if first:
            return True
        second = self._run_entailment_check(request)
        return bool(second)

    def _run_entailment_check(self, request) -> bool | None:
        """One entailment call. `None` = outage/malformed (fails closed by the
        caller without a retry); `True`/`False` = the judge's verdict."""
        try:
            raw = self._generator.generate(request.system, request.user, request.schema)
        except AnswerGenerationUnavailable:
            return None
        try:
            payload = json.loads(raw)
            entailed = payload["entailed"]
            unsupported = payload["unsupported"]
        except (ValueError, TypeError, KeyError):
            return None
        if not isinstance(entailed, bool) or not isinstance(unsupported, list):
            return None
        return entailed and not unsupported

    def _check_sufficiency(
        self, query: str, evidence_texts: tuple[str, ...], intro: bool = False
    ) -> str | None:
        """A focused reasoning call: is the turn specific enough to answer, or
        must we ask? Returns a clarifying question, or None to proceed.

        Skipped without a model, for a bare-name intro (not a dose), or for a
        single evidence block (nothing to disambiguate)."""
        if self._generator is None or intro or len(evidence_texts) < 2:
            return None
        request = build_sufficiency_request(query, evidence_texts)
        try:
            raw = self._generator.generate(request.system, request.user, request.schema)
        except AnswerGenerationUnavailable:
            return None
        try:
            payload = json.loads(raw)
        except (ValueError, TypeError):
            return None
        if isinstance(payload, dict) and payload.get("sufficient") is False:
            question = payload.get("clarifying_question")
            if isinstance(question, str) and question.strip():
                return question.strip()
        return None

    @staticmethod
    def _cited_only(
        indexed: list[tuple[int, Citation]], answer_text: str
    ) -> tuple[Citation, ...]:
        """Keep citations whose 1-based evidence marker [n] appears in the text."""
        used = {int(m) for m in re.findall(r"\[(\d+)\]", answer_text)}
        return tuple(citation for index, citation in indexed if index in used)

    @staticmethod
    def _indexed_citations(
        result: RetrievalResult,
    ) -> list[tuple[int, Citation]] | None:
        """Citations tagged with the 1-based evidence index the prompt gives them,
        so the response can show only the ones the answer cited."""
        citations: list[tuple[int, Citation]] = []
        for index, evidence in enumerate(result.evidence, start=1):
            if not evidence.source_refs:
                return None
            for source in evidence.source_refs:
                printed_range = source.printed_page_range
                if printed_range is not None:
                    start, end = printed_range
                elif source.printed_page is not None:
                    start = end = source.printed_page
                else:
                    return None
                citations.append((index, Citation(
                    chunk_id=evidence.matched_doc_id,
                    printed_page_start=int(start),
                    printed_page_end=int(end),
                    physical_page=source.physical_page,
                    block_id=source.block_id,
                    bbox=source.bbox,
                    source_crop=source.source_crop,
                    # Backward-compatible compact attachment identifier. A
                    # real crop path wins; otherwise the block id plus the
                    # structured page/bbox fields is enough to render later.
                    attachment=source.source_crop or source.block_id,
                )))
        return citations