from __future__ import annotations

import uuid
from typing import Annotated, Any, Protocol

from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, Field

from rag.answer import GroundedAnswerService
from rag.metrics import TRACE_WRITE_FAILED, Metrics, NullMetrics
from rag.models import QueryIntent, SubjectScope
from rag.policy import resolve_subject_scope


class TraceWriter(Protocol):
    def save(self, **fields: Any) -> str: ...


class RagQueryRequest(BaseModel):
    query: str = Field(min_length=1, max_length=4000)
    subject_scope: SubjectScope
    intent: QueryIntent
    # Optional: when present, the turn is answered in conversation context
    # (follow-up inheritance, clarify, smalltalk). Absent → single-turn, exactly
    # as before, so existing callers are unchanged.
    conversation_id: str | None = Field(default=None, max_length=128)


class CitationResponse(BaseModel):
    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 chunk text this citation stands for — lets the UI
    # show precisely what was retrieved, not a client-side guess at it.
    evidence_text: str = ""


class AnswerClaimResponse(BaseModel):
    text: str
    source_ids: list[str]


class AnswerBlockResponse(BaseModel):
    title: str
    kind: str
    claims: list[AnswerClaimResponse]


class AnswerPlanResponse(BaseModel):
    verbosity: str
    layout: str
    reasoning_mode: str
    show_heading: bool
    needs_warning: bool


class RagQueryResponse(BaseModel):
    trace_id: str
    decision: str
    reason: str
    answer: str | None
    resolved_drug_id: str | None
    citations: list[CitationResponse]
    # Whether `answer` is an LLM paraphrase (verified by grounding +
    # entailment) or a verbatim extractive quote of the retrieved source —
    # the UI shows these differently so a clinician knows which they're
    # reading.
    generated: bool = False
    # Short suggested replies for a `decision == "clarify"` turn (e.g.
    # ["Người lớn", "Trẻ em"]) — only populated by the sufficiency-check
    # clarify path today; other clarify sources (no_drug, dosing_calc's
    # needs_clarify) leave this empty rather than fabricate options.
    quick_replies: list[str] = []
    blocks: list[AnswerBlockResponse] = []
    answer_mode: str = "concise"
    answer_plan: AnswerPlanResponse | None = None


def _answer_service(request: Request) -> GroundedAnswerService:
    service = getattr(request.app.state, "answer_service", None)
    if service is None:
        raise HTTPException(status_code=503, detail="RAG backend is not configured")
    return service


def _trace_writer(request: Request) -> TraceWriter:
    writer = getattr(request.app.state, "trace_writer", None)
    if writer is None:
        raise HTTPException(status_code=503, detail="Trace database is not configured")
    return writer


def _metrics(request: Request) -> Metrics:
    return getattr(request.app.state, "metrics", None) or NullMetrics()


router = APIRouter(prefix="/v1/rag", tags=["rag"])


class SuggestResponse(BaseModel):
    suggestions: list[str]


@router.get("/suggest", response_model=SuggestResponse)
def suggest_drugs(q: str, request: Request) -> SuggestResponse:
    """As-you-type drug-name autocomplete, so a name is picked, not mistyped."""
    agent = getattr(request.app.state, "conversational", None)
    if agent is None or not q.strip():
        return SuggestResponse(suggestions=[])
    return SuggestResponse(suggestions=agent.complete(q.strip()))


def _map_citations(items) -> list[CitationResponse]:
    return [
        CitationResponse(
            chunk_id=item.chunk_id,
            printed_page_start=item.printed_page_start,
            printed_page_end=item.printed_page_end,
            physical_page=item.physical_page,
            block_id=item.block_id,
            bbox=item.bbox,
            source_crop=item.source_crop,
            attachment=item.attachment,
            evidence_text=item.evidence_text,
        )
        for item in items
    ]


def _map_blocks(items) -> list[AnswerBlockResponse]:
    return [
        AnswerBlockResponse(
            title=item.title,
            kind=item.kind,
            claims=[
                AnswerClaimResponse(text=claim.text, source_ids=list(claim.source_ids))
                for claim in item.claims
            ],
        )
        for item in items
    ]


def _map_plan(item) -> AnswerPlanResponse | None:
    if item is None:
        return None
    return AnswerPlanResponse(
        verbosity=item.verbosity,
        layout=item.layout,
        reasoning_mode=item.reasoning_mode,
        show_heading=item.show_heading,
        needs_warning=item.needs_warning,
    )


@router.post("/query", response_model=RagQueryResponse)
def query_rag(
    payload: RagQueryRequest,
    request: Request,
    answers: Annotated[GroundedAnswerService, Depends(_answer_service)],
    traces: Annotated[TraceWriter, Depends(_trace_writer)],
    metrics: Annotated[Metrics, Depends(_metrics)],
) -> RagQueryResponse:
    agent = getattr(request.app.state, "conversational", None)

    # `payload.subject_scope`/`payload.intent` are what the CALLER claims —
    # logged below for audit, but the RagAgent path does not take them as an
    # input at all. It derives scope from the query text itself (the same
    # `rag.policy` check used here for the no-agent fallback) as part of its
    # own LLM understanding call, and does not gate on intent at all (this
    # product is for doctors/pharmacists; a client label must not be, and
    # here structurally cannot be, the safety decision — F-02).
    subject_scope = resolve_subject_scope(payload.query, payload.subject_scope)
    intent = payload.intent

    if agent is not None:
        # The live path (F-03): one LLM call understands the turn (drug
        # identity against the real catalog, turn type, population/weight),
        # then routes to the safety-verified retrieval + grounded-answer
        # engine. Replaces the old fuzzy resolver + keyword section router +
        # manual follow-up inheritance for both single- and multi-turn.
        reply = agent.handle(payload.query, payload.conversation_id)
        decision = reply.decision
        reason = reply.reason
        answer = reply.clarification if reply.clarification is not None else reply.answer
        resolved_drug_id = ", ".join(reply.drugs) if reply.drugs else None
        citations = _map_citations(reply.citations)
        generated = reply.generated
        quick_replies = list(reply.quick_replies)
        blocks = _map_blocks(reply.blocks)
        answer_mode = reply.answer_mode
        answer_plan = _map_plan(reply.plan)
    else:
        # No generator configured (ANSWER_PROVIDER=disabled): there is no LLM
        # to understand a turn with, so this is retrieval-only, single-turn,
        # unchanged from before F-03.
        grounded = answers.answer(payload.query, subject_scope, intent)
        if grounded.clarification is not None:
            decision, reason = "clarify", "needs_more_info"
            answer = grounded.clarification
            resolved_drug_id = grounded.result.resolved_drug_id
            citations = []
            generated = False
            quick_replies = list(grounded.quick_replies)
            blocks = []
            answer_mode = "concise"
            answer_plan = None
        else:
            decision = grounded.result.decision.value
            reason = grounded.result.reason
            answer = grounded.answer
            resolved_drug_id = grounded.result.resolved_drug_id
            citations = _map_citations(grounded.citations)
            generated = grounded.generated
            quick_replies = []
            blocks = _map_blocks(grounded.blocks)
            answer_mode = grounded.answer_mode
            answer_plan = _map_plan(grounded.plan)

    # Trace persistence is fail-open (F-09): an already-computed, safe answer
    # must reach the caller even if Postgres is unreachable. `save()` opens a
    # fresh connection per call with no retry, so a DB outage previously
    # turned a good answer into a 500 for a reason that has nothing to do
    # with whether the answer was safe. `trace_id` degrades to a local,
    # unpersisted uuid — still a valid response field, just not one `GET
    # /v1/rag/trace/{id}` (if it existed) could later look up.
    try:
        trace_id = traces.save(
            query=payload.query,
            # The resolved (server-derived) values, not the caller's claim —
            # this is what actually gated the answer, so it's what the trace
            # must show.
            subject_scope=subject_scope.value,
            intent=intent.value,
            decision=decision,
            reason=reason,
            resolved_drug_id=resolved_drug_id,
            citations=tuple(item.model_dump() for item in citations),
        )
    except Exception:
        metrics.increment(TRACE_WRITE_FAILED)
        trace_id = str(uuid.uuid4())
    return RagQueryResponse(
        trace_id=trace_id,
        decision=decision,
        reason=reason,
        answer=answer,
        resolved_drug_id=resolved_drug_id,
        citations=citations,
        generated=generated,
        quick_replies=quick_replies,
        blocks=blocks,
        answer_mode=answer_mode,
        answer_plan=answer_plan,
    )