from __future__ import annotations

import uuid
from typing import Annotated, Any, Literal, Protocol

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

from rag.answer import DISCLAIMER, GroundedAnswerService
from rag.metrics import DECISION, TRACE_WRITE_FAILED, Metrics, NullMetrics
from rag.models import QueryIntent, SubjectScope
from rag.policy import resolve_subject_scope
from rag.telemetry import (
    annotate_current_span,
    current_correlation_id,
    current_trace_id,
    stage,
)
from adapters.postgres import FeedbackTraceNotFound


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

    def save_feedback(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 = ""
    drug_id: str | None = None
    drug_name: str | None = None
    section_key: str | None = None
    section_title: str | None = None
    source_document: str = "Dược thư Quốc gia Việt Nam 2018"


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 MedicationCandidateAssessmentResponse(BaseModel):
    drug_id: str
    drug_name: str
    indication_supported: bool
    status: str
    indication_source_ids: list[str]
    safety_source_ids: list[str]


class RagFeedbackRequest(BaseModel):
    trace_id: uuid.UUID
    rating: Literal["helpful", "not_helpful"]
    comment: str | None = Field(default=None, max_length=2000)
    conversation_id: str | None = Field(default=None, max_length=128)


class RagFeedbackResponse(BaseModel):
    feedback_id: str
    status: Literal["saved"] = "saved"


class RagQueryResponse(BaseModel):
    trace_id: str
    correlation_id: str
    otel_trace_id: str | None = None
    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
    candidate_assessments: list[MedicationCandidateAssessmentResponse] = []
    # `docs/architecture.md` specifies the disclaimer at several layers. The
    # web banner was the only one in place, so any other consumer received
    # medical content with nothing attached. It is a fixed string owned by
    # `rag/answer.py`, never produced by the model, and it is present on
    # every decision — an abstain or a clarification is still a clinical
    # response. Defaulted here as well so a response constructed in a test or
    # a future code path cannot accidentally omit it.
    disclaimer: str = DISCLAIMER


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"])


@router.post("/feedback", response_model=RagFeedbackResponse)
def save_feedback(
    payload: RagFeedbackRequest,
    traces: Annotated[TraceWriter, Depends(_trace_writer)],
) -> RagFeedbackResponse:
    comment = payload.comment.strip() if payload.comment else None
    try:
        feedback_id = traces.save_feedback(
            trace_id=str(payload.trace_id),
            rating=payload.rating,
            comment=comment or None,
            conversation_id=payload.conversation_id,
        )
    except FeedbackTraceNotFound as exc:
        raise HTTPException(status_code=404, detail="trace_not_found") from exc
    except Exception as exc:
        raise HTTPException(status_code=503, detail="feedback_store_unavailable") from exc
    return RagFeedbackResponse(feedback_id=feedback_id)


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,
            drug_id=item.drug_id,
            drug_name=item.drug_name,
            section_key=item.section_key,
            section_title=item.section_title,
            source_document=item.source_document,
        )
        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,
    )


def _map_candidate_assessments(items) -> list[MedicationCandidateAssessmentResponse]:
    return [
        MedicationCandidateAssessmentResponse(
            drug_id=item.drug_id,
            drug_name=item.drug_name,
            indication_supported=item.indication_supported,
            status=item.status.value,
            indication_source_ids=[
                evidence.matched_doc_id for evidence in item.indication_evidence
            ],
            safety_source_ids=[
                evidence.matched_doc_id
                for evidence in item.evidence
                if evidence not in item.indication_evidence
            ],
        )
        for item in items
    ]


@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:
    with stage("receive"):
        agent = getattr(request.app.state, "conversational", None)
        subject_scope = resolve_subject_scope(payload.query, payload.subject_scope)
        intent = payload.intent

    # `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).
    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)
        candidate_assessments = _map_candidate_assessments(
            reply.candidate_assessments
        )
    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.
        with stage("routing"):
            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
            candidate_assessments = []
        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)
            candidate_assessments = []

    # 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.
    correlation_id = current_correlation_id()
    otel_trace_id = current_trace_id()
    try:
        with stage("persistence"):
            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),
                correlation_id=correlation_id,
                otel_trace_id=otel_trace_id,
            )
    except Exception:
        metrics.increment(TRACE_WRITE_FAILED)
        trace_id = str(uuid.uuid4())
    metrics.increment(DECISION, decision=decision, reason=reason)
    annotate_current_span(
        **{
            "duocthu.decision": decision,
            "duocthu.reason": reason,
            "duocthu.citation_count": len(citations),
            "duocthu.generated": generated,
            "duocthu.persisted_trace_id": trace_id,
        }
    )
    with stage("response"):
        response = RagQueryResponse(
            trace_id=trace_id,
            correlation_id=correlation_id,
            otel_trace_id=otel_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,
            candidate_assessments=candidate_assessments,
        )
    return response