"""Two answer-UX fixes, pinned:

- citations shown = only the sources the answer cited, not every retrieved chunk;
- a bare drug name is introduced, not restated section-by-section.
"""
from __future__ import annotations

import json

from rag.answer import GroundedAnswerService
from rag.models import (
    Evidence,
    EvidenceDecision,
    QueryIntent,
    RetrievalResult,
    SourceRef,
    SubjectScope,
)
from rag.prompt import build_request


def _evidence(i: int, page: int) -> Evidence:
    return Evidence(
        evidence_id=f"drug::sec::{i}",
        matched_doc_id=f"drug::sec::{i}",
        kind="prose",
        text=f"đoạn bằng chứng {i}",
        score=1.0,
        source_refs=(SourceRef(physical_page=page, precision="exact", printed_page=page),),
        hydrated_from_parent=False,
        requires_visual_check=False,
    )


class _Routing:
    def __init__(self, result: RetrievalResult) -> None:
        self._result = result

    def retrieve(self, query, subject_scope, intent):  # noqa: ARG002
        return self._result


class _Generator:
    def __init__(self, payload: dict) -> None:
        self._payload = payload

    def generate(self, system: str, user: str, schema: dict) -> str:  # noqa: ARG002
        return json.dumps(self._payload, ensure_ascii=False)


def _answerable(*evidence: Evidence, is_overview: bool = False) -> RetrievalResult:
    return RetrievalResult(
        EvidenceDecision.ANSWERABLE,
        "grounded_evidence_available",
        tuple(evidence),
        resolved_drug_id="drug",
        is_drug_overview=is_overview,
    )


def test_only_cited_sources_are_returned():
    result = _answerable(_evidence(0, 100), _evidence(1, 200), _evidence(2, 300))
    service = GroundedAnswerService(
        _Routing(result),
        _Generator({"answer": "Chỉ dùng đoạn hai [2].", "evidence_sufficient": True}),
    )

    grounded = service.answer("q", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)

    assert grounded.generated is True
    assert len(grounded.citations) == 1
    assert grounded.citations[0].printed_page_start == 200


def test_answer_citing_nothing_falls_back_to_all_citations():
    result = _answerable(_evidence(0, 100), _evidence(1, 200))
    service = GroundedAnswerService(
        _Routing(result),
        # no [n] marker at all: rather than show zero provenance, show all.
        _Generator({"answer": "Không có trích dẫn.", "evidence_sufficient": True}),
    )

    grounded = service.answer("q", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)

    assert len(grounded.citations) == 2


def test_bare_name_builds_an_intro_prompt():
    intro = build_request("PARACETAMOL", ("đoạn A", "đoạn B"), intro=True)
    assert "GIỚI THIỆU" in intro.user
    assert "CÂU HỎI:" not in intro.user

    normal = build_request("Liều?", ("đoạn A",), intro=False)
    assert "CÂU HỎI:" in normal.user
    assert "GIỚI THIỆU" not in normal.user