from __future__ import annotations

import json
import re
from dataclasses import dataclass, replace

from . import grounding, metrics as metric_names
from .budget import RequestBudget, RequestBudgetExhausted
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

# See `_verify_entailment`'s docstring for the measured trade-off behind
# widening this from 2 to 3.
_ENTAILMENT_MAX_ATTEMPTS = 3


@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
    # The exact retrieved text this citation stands for — the same string
    # handed to the generator/entailment checks, so the UI can show precisely
    # what was retrieved rather than a fabricated summary of it.
    evidence_text: str = ""


@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
    # Short suggested replies for `clarification`, e.g. ("Người lớn", "Trẻ
    # em") — only populated when the sufficiency check judged the question
    # to have a few natural discrete answers, never invented client-side.
    quick_replies: tuple[str, ...] = ()


@dataclass(frozen=True)
class _GenOutcome:
    answer: str | None = None
    clarification: str | None = None
    # The specific reason a rejection happened — the exact string already
    # used for the GENERATION_REJECTED metric, propagated here so
    # `answer_from_result` can put it in the API response's `reason` field
    # instead of a generic catch-all. `None` when `answer`/`clarification`
    # is set (nothing was rejected).
    reject_reason: str | None = None


@dataclass(frozen=True)
class _RawAttempt:
    """One raw `_attempt_generation` call, before any metric is charged —
    lets `_generate` retry the noisy `insufficient` case without
    double-counting a rejection metric across both attempts."""
    answer: str | None = None
    clarification: str | None = None
    insufficient: bool = False
    outage: bool = False
    budget_exhausted: bool = False
    malformed: bool = False


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** with the specific reason that failed
      it (`provider_unavailable`, `malformed_output`, `evidence_insufficient`,
      a `grounding.verify` reason, or `unsupported_claim`; falls back to
      the generic `generation_unavailable` only if none of those was set)
      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, list_mode: bool = False,
        budget: RequestBudget | None = None,
    ) -> 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.

        `list_mode=True`: the evidence is several DIFFERENT drugs' own
        sections (symptom_to_drug), not alternative phrasings of one drug's
        answer — the sufficiency clarify ("which kind of headache?") that's
        right for a single dose question doesn't fit a reverse lookup, whose
        whole point is to show what the formulary has and let the clinician
        narrow it themselves; skipped here the same way a bare-name intro
        already skips it.

        `budget` (F-08): threaded through to every LLM call this method
        makes (sufficiency, generate, up to 2 entailment). `None` (the
        default) means unbounded, unchanged from before F-08 — only
        `RagAgent` constructs a real budget today.
        """
        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.
        sufficiency = (
            None if list_mode
            else self._check_sufficiency(
                query, evidence_texts, result.is_drug_overview, budget=budget
            )
        )
        if sufficiency is not None:
            clarify_q, quick_replies = sufficiency
            return GroundedAnswer(
                result, clarify_q, (), clarification=clarify_q, quick_replies=quick_replies
            )

        outcome = self._generate(
            query, evidence_texts, intro=result.is_drug_overview, list_mode=list_mode,
            budget=budget,
        )
        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.
            # The specific check that failed (provider_unavailable,
            # malformed_output, evidence_insufficient, ungrounded_number,
            # uncited_claim, unsupported_claim, request_budget_exhausted) —
            # found live 2026-08-07: every one of these used to collapse into
            # the same generic "generation_unavailable" by the time it
            # reached the API response/trace, so a real, diagnosable cause
            # (e.g. a genuine provider outage) was indistinguishable from
            # ordinary entailment noise without reading server-side metrics
            # by hand. `outcome.reject_reason` already carries the granular
            # value the metric above uses — just propagate it.
            reason = outcome.reject_reason or "generation_unavailable"
            self._metrics.increment(metric_names.ABSTENTION, reason=reason)
            return GroundedAnswer(
                replace(
                    result,
                    decision=EvidenceDecision.ABSTAIN,
                    reason=reason,
                ),
                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 _attempt_generation(
        self, request: "GenerationRequest", budget: RequestBudget | None
    ) -> "_RawAttempt":
        """One raw generation call, parsed but not yet metric-counted or
        verified — the caller decides whether to retry before charging a
        metric to any particular reason."""
        try:
            if budget is not None:
                budget.require()
            raw = self._generator.generate(request.system, request.user, request.schema)
        except RequestBudgetExhausted:
            return _RawAttempt(budget_exhausted=True)
        except AnswerGenerationUnavailable:
            return _RawAttempt(outage=True)

        try:
            payload = json.loads(raw)
            answer = payload["answer"]
            sufficient = payload["evidence_sufficient"]
        except (ValueError, TypeError, KeyError):
            return _RawAttempt(malformed=True)

        # 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 _RawAttempt(clarification=clarify.strip())

        if not isinstance(answer, str) or not isinstance(sufficient, bool):
            return _RawAttempt(malformed=True)
        if not sufficient:
            return _RawAttempt(insufficient=True)
        return _RawAttempt(answer=answer)

    def _generate(
        self, query: str, evidence_texts: tuple[str, ...], intro: bool = False,
        list_mode: bool = False, budget: RequestBudget | None = None,
    ) -> "_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, list_mode=list_mode)
        attempt = self._attempt_generation(request, budget)
        if attempt.insufficient:
            # Empirically noisy (found live 2026-08-07, reproduced 3/3 on a
            # fresh retry): the model's own evidence_sufficient=false
            # self-assessment sometimes flips to a correct, fully grounded
            # answer when asked again with the IDENTICAL evidence — the same
            # one-retry pattern `_verify_entailment` already uses below for
            # its own noisy judge call. Only the terminal "insufficient AND
            # no clarifying question" case retries; a legitimate ask-for-
            # more-detail clarify is untouched.
            attempt = self._attempt_generation(request, budget)

        if attempt.budget_exhausted:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="request_budget_exhausted"
            )
            return _GenOutcome(reject_reason="request_budget_exhausted")
        if attempt.outage:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="provider_unavailable"
            )
            return _GenOutcome(reject_reason="provider_unavailable")
        if attempt.malformed:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="malformed_output"
            )
            return _GenOutcome(reject_reason="malformed_output")
        if attempt.clarification is not None:
            return _GenOutcome(clarification=attempt.clarification)
        if attempt.insufficient:
            # The model says the evidence does not answer the question, on
            # both attempts. Showing the retrieved section verbatim lets the
            # clinician judge that.
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="evidence_insufficient"
            )
            return _GenOutcome(reject_reason="evidence_insufficient")

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

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

    def _verify_entailment(
        self, answer: str, evidence_texts: tuple[str, ...],
        budget: RequestBudget | None = None,
    ) -> 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 — that failure mode is
        reliable, not noisy, so it stops immediately rather than spending
        retries on it. A single rejection is NOT reliable: 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. Up to `_ENTAILMENT_MAX_ATTEMPTS` same-claim calls run;
        accept on the first `True`, discard only if every attempt agrees
        reject. Widened from 2 to 3 attempts 2026-08-07 after a live
        adversarial sample (50 real questions) measured this specific check
        as roughly half of all false abstentions on genuinely answerable
        questions. Trade-off, stated plainly: this raises the bar a
        genuinely fabricated claim must now clear too (it survives if ANY
        one of 3 noisy calls wrongly accepts it, not just 1 of 2) — accepted
        because the probed noise is symmetric and the entailment prompt
        itself is unchanged, not because the risk is zero.
        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)
        for _ in range(_ENTAILMENT_MAX_ATTEMPTS):
            verdict = self._run_entailment_check(request, budget=budget)
            if verdict is None:
                return False
            if verdict:
                return True
        return False

    def _run_entailment_check(
        self, request, budget: RequestBudget | None = None
    ) -> bool | None:
        """One entailment call. `None` = outage/malformed/budget-exhausted
        (fails closed by the caller without a retry); `True`/`False` = the
        judge's verdict."""
        try:
            if budget is not None:
                budget.require()
            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,
        budget: RequestBudget | None = None,
    ) -> tuple[str, tuple[str, ...]] | None:
        """A focused reasoning call: is the turn specific enough to answer, or
        must we ask? Returns (clarifying_question, quick_replies), or None to
        proceed. `quick_replies` is often empty — only populated when the
        model judged the missing detail has a few natural discrete answers
        (e.g. "Người lớn"/"Trẻ em"), never invented here.

        Skipped without a model, for a bare-name intro (not a dose), or for a
        single evidence block (nothing to disambiguate).

        Fails OPEN on outage/budget-exhaustion (returns None, proceeds to
        generate) — deliberately different from every other call in this
        file, which fail closed. This is a reasoning heuristic, not a safety
        check; grounding + entailment remain the real gate on whatever gets
        generated next, so skipping this one costs UX quality (a dose
        question that should have asked for age/weight might not), not
        safety."""
        if self._generator is None or intro or len(evidence_texts) < 2:
            return None
        request = build_sufficiency_request(query, evidence_texts)
        try:
            if budget is not None:
                budget.require()
            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():
                raw_replies = payload.get("quick_replies")
                replies = tuple(
                    reply.strip()
                    for reply in raw_replies
                    if isinstance(reply, str) and reply.strip()
                ) if isinstance(raw_replies, list) else ()
                return question.strip(), replies
        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,
                    evidence_text=evidence.text,
                )))
        return citations