"""The bounded reasoning loop.

Understand → plan → retrieve → assess → refine → generate → verify → repair.
Every edge is bounded, and every budget is decremented **before** the call it
pays for, so exhaustion degrades to the best answer so far rather than to an
error.

Two rules hold across every path and are the reason this can be added to a
formulary at all:

- `grounding.verify` still gates every generated answer. Reasoning chooses what
  to look up and how to phrase it; it is never a source of facts.
- A clarify signal bypasses the loop entirely. Asking beats guessing, and the
  signals are resolver states — ambiguous drug, unresolved attribute — not a
  model's confidence score.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
from typing import Protocol

from . import metrics as metric_names
from .conversation import ConversationState, ResolvedQuestion
from .metrics import Metrics, NullMetrics

MAX_RETRIEVAL_ROUNDS = 2
MAX_REPAIRS = 1
MAX_LLM_CALLS = 4
MAX_WALL_CLOCK_MS = 20_000


class BudgetExhausted(RuntimeError):
    """Raised only inside the loop, never surfaced; the loop catches it."""


@dataclass
class TurnBudget:
    """Mutable on purpose: one budget is threaded through one turn."""

    llm_calls: int = MAX_LLM_CALLS
    retrieval_rounds: int = MAX_RETRIEVAL_ROUNDS
    repairs: int = MAX_REPAIRS
    wall_clock_ms: int = MAX_WALL_CLOCK_MS
    elapsed_ms: int = 0

    def spend_llm(self) -> None:
        if self.llm_calls <= 0:
            raise BudgetExhausted("llm_calls")
        self.llm_calls -= 1

    def spend_retrieval(self) -> None:
        if self.retrieval_rounds <= 0:
            raise BudgetExhausted("retrieval_rounds")
        self.retrieval_rounds -= 1

    def spend_repair(self) -> None:
        if self.repairs <= 0:
            raise BudgetExhausted("repairs")
        self.repairs -= 1

    def out_of_time(self) -> bool:
        return self.elapsed_ms >= self.wall_clock_ms


class ClarifyReason:
    AMBIGUOUS_DRUG = "ambiguous_drug"
    NO_ATTRIBUTE = "no_attribute"
    MULTI_ATTRIBUTE = "multi_attribute"
    STILL_INSUFFICIENT = "still_insufficient"


@dataclass(frozen=True)
class Clarification:
    reason: str
    question: str
    options: tuple[str, ...] = ()


@dataclass(frozen=True)
class Sufficiency:
    """The assessor's verdict on retrieved evidence.

    `missing` must name something specific — a section, a population, a second
    drug. "Feels incomplete" does not buy a retrieval round; a round is only
    spent when there is a concrete thing to go and fetch.
    """

    sufficient: bool
    missing: str | None = None
    refined_query: str | None = None


class SufficiencyAssessor(Protocol):
    def assess(
        self, resolved: ResolvedQuestion, evidence_texts: tuple[str, ...]
    ) -> Sufficiency: ...


class DeterministicAssessor:
    """The no-LLM default, and the reference for what the port must do.

    Runs offline and is what the loop uses until a provider is enabled. It only
    reports insufficiency it can *demonstrate* — a population was asked for and
    no retrieved text mentions it — so it can never spin the loop on a feeling.
    """

    POPULATION_TERMS = {
        "nguoi_lon": ("người lớn",),
        "tre_em": ("trẻ em", "trẻ nhỏ", "trẻ "),
        "tre_so_sinh": ("sơ sinh",),
        "phu_nu_co_thai": ("thai", "mang thai"),
        "phu_nu_cho_con_bu": ("cho con bú", "sữa mẹ"),
        "nguoi_cao_tuoi": ("người cao tuổi", "người già"),
        "suy_than": ("suy thận", "clcr"),
        "suy_gan": ("suy gan",),
    }

    def assess(
        self, resolved: ResolvedQuestion, evidence_texts: tuple[str, ...]
    ) -> Sufficiency:
        if not evidence_texts:
            return Sufficiency(False, missing="no_evidence")
        if resolved.population is None:
            return Sufficiency(True)

        terms = self.POPULATION_TERMS.get(resolved.population, ())
        haystack = " ".join(evidence_texts).casefold()
        if any(term in haystack for term in terms):
            return Sufficiency(True)
        return Sufficiency(
            False,
            missing=f"population:{resolved.population}",
            refined_query=f"{resolved.text} {terms[0] if terms else ''}".strip(),
        )


@dataclass(frozen=True)
class LoopOutcome:
    """What one turn produced, plus what it cost."""

    answer: str | None
    clarification: Clarification | None
    evidence_texts: tuple[str, ...]
    retrieval_rounds_used: int
    repairs_used: int
    stopped_because: str
    generated: bool = False


@dataclass
class LoopTrace:
    """Ordered record of stages, for the dashboard and for debugging."""

    stages: list[str] = field(default_factory=list)

    def enter(self, stage: str) -> None:
        self.stages.append(stage)


def clarify_for(
    reason: str, options: tuple[str, ...] = ()
) -> Clarification:
    questions = {
        ClarifyReason.NO_ATTRIBUTE: (
            "Anh/chị muốn tra thuộc tính nào của thuốc này?"
        ),
        ClarifyReason.AMBIGUOUS_DRUG: (
            "Câu hỏi có thể ứng với nhiều thuốc. Anh/chị muốn tra thuốc nào?"
        ),
        ClarifyReason.MULTI_ATTRIBUTE: (
            "Câu hỏi nhắc tới nhiều mục. Anh/chị muốn xem mục nào trước?"
        ),
        ClarifyReason.STILL_INSUFFICIENT: (
            "Chưa tìm đủ căn cứ trong Dược thư cho ý này. "
            "Anh/chị có thể nêu rõ hơn điều cần tra không?"
        ),
    }
    return Clarification(reason, questions[reason], options)


class Retrieve(Protocol):
    def __call__(self, resolved: ResolvedQuestion) -> tuple[str, ...]: ...


class Generate(Protocol):
    def __call__(
        self, resolved: ResolvedQuestion, evidence: tuple[str, ...], state: ConversationState
    ) -> str | None: ...


def run_turn(
    state: ConversationState,
    resolved: ResolvedQuestion,
    retrieve: Retrieve,
    generate: Generate,
    clarify_signals: tuple[str, ...] = (),
    assessor: SufficiencyAssessor | None = None,
    budget: TurnBudget | None = None,
    metrics: Metrics | None = None,
    trace: LoopTrace | None = None,
) -> LoopOutcome:
    """One conversational turn through the bounded loop.

    `clarify_signals` comes from the existing resolvers — ambiguous drug,
    unresolved section, multi-attribute. They short-circuit before any spend,
    because a question worth asking is cheaper and safer than a guess.
    """
    budget = budget or TurnBudget()
    assessor = assessor or DeterministicAssessor()
    metrics = metrics or NullMetrics()
    trace = trace or LoopTrace()

    trace.enter("understand")
    if clarify_signals:
        reason = clarify_signals[0]
        metrics.increment(metric_names.CLARIFY_ASKED, reason=reason)
        trace.enter("clarify")
        return LoopOutcome(
            answer=None,
            clarification=clarify_for(reason),
            evidence_texts=(),
            retrieval_rounds_used=0,
            repairs_used=0,
            stopped_because="clarify_signal",
        )

    evidence: tuple[str, ...] = ()
    rounds_used = 0
    stopped = "sufficient"

    while True:
        try:
            budget.spend_retrieval()
        except BudgetExhausted:
            stopped = "retrieval_budget"
            break
        trace.enter("retrieve")
        evidence = retrieve(resolved)
        rounds_used += 1

        trace.enter("assess")
        verdict = assessor.assess(resolved, evidence)
        if verdict.sufficient:
            break
        if budget.retrieval_rounds <= 0 or budget.out_of_time():
            stopped = "retrieval_budget"
            break
        # A round is spent only on a named gap with a genuinely new query.
        if not verdict.missing or not verdict.refined_query:
            stopped = "no_actionable_gap"
            break
        if verdict.refined_query == resolved.text:
            stopped = "query_unchanged"
            break
        trace.enter("refine")
        metrics.increment(metric_names.LOOP_REFINED, missing=verdict.missing)
        resolved = replace(resolved, text=verdict.refined_query)

    metrics.increment(metric_names.LOOP_ROUNDS, rounds=str(rounds_used))

    if not evidence:
        trace.enter("clarify")
        metrics.increment(
            metric_names.CLARIFY_ASKED, reason=ClarifyReason.STILL_INSUFFICIENT
        )
        return LoopOutcome(
            answer=None,
            clarification=clarify_for(ClarifyReason.STILL_INSUFFICIENT),
            evidence_texts=(),
            retrieval_rounds_used=rounds_used,
            repairs_used=0,
            stopped_because="no_evidence",
        )

    repairs_used = 0
    answer: str | None = None
    while True:
        trace.enter("generate")
        try:
            budget.spend_llm()
        except BudgetExhausted:
            stopped = "llm_budget"
            break
        answer = generate(resolved, evidence, state)
        if answer is not None:
            break
        # `generate` returning None means verification already refused it.
        try:
            budget.spend_repair()
        except BudgetExhausted:
            stopped = "repair_budget"
            break
        repairs_used += 1
        trace.enter("repair")
        metrics.increment(metric_names.LOOP_REPAIRED)

    return LoopOutcome(
        answer=answer,
        clarification=None,
        evidence_texts=evidence,
        retrieval_rounds_used=rounds_used,
        repairs_used=repairs_used,
        stopped_because=stopped,
        generated=answer is not None,
    )