"""RAG knowledge base (#60): corpus integrity, retrieval quality, grounded and safe answers."""

import re

import pytest

from app.core.types import ScrubbedText
from app.llm.gateway import LLMGateway
from app.llm.mock import MockTasks
from app.llm.rag import qa as answer_module
from app.llm.rag.corpus import heading, load_chunks, load_config, rag_dir, split_sentences
from app.llm.rag.index import get_index, tokens
from app.llm.rag.qa import LIT_REFS, PRESCRIBING, QUOTE_FILTER, answer, fallback_text, select
from app.llm.rag.stemmer import stem

S = ScrubbedText

# question -> (document, pattern that must occur in one of the top-3 chunks). Answers checked by hand
# against the PDFs; the set doubles as the tuning set, so the threshold below guards against regressions.
GOLD = [
    ("через сколько начинают действовать антидепрессанты", "kr_depression", r"Тимоаналептический эффект обычно"),
    ("когда антидепрессант даст эффект", "kr_depression", r"Тимоаналептический эффект обычно"),
    ("помогает ли ДПДГ при ПТСР", "kr_ptsd", r"ДПДГ"),
    ("какая психотерапия эффективна при посттравматическом расстройстве", "kr_ptsd", r"наибольшей доказательной"),
    ("как правильно отменять антидепрессанты", "kr_depression", r"постепенное снижение дозы"),
    ("что такое флэшбек", "kr_ptsd", r"Флэшбек \(flashback"),
    ("шкала Гамильтона для депрессии", "kr_depression", r"Шкала Гамильтона для оценки депрессии"),
    ("помощь военнослужащим после боевой психической травмы", "kr_ptsd", r"военнослужащ"),
    ("колумбийская шкала суицидального риска", "kr_depression", r"Колумбийская шкала"),
    ("вызывают ли бензодиазепины зависимость", "kr_ptsd", r"формировать зависимость"),
    ("что сказать пациенту о ПТСР", "kr_ptsd", r"Что такое посттравматическое"),
    ("электросудорожная терапия при резистентной депрессии", "kr_depression", r"лектросудорожн"),
    ("скрининговый опросник ПТСР", "kr_ptsd", r"Trauma Screening"),
    ("физическая активность при депрессии", "kr_depression", r"физическ"),
    ("сертралин при ПТСР", "kr_ptsd", r"ертралин"),
    ("противорецидивная терапия депрессии сколько длится", "kr_depression", r"ротиворецидивн"),
    ("правило M7 срочные симптомы", "meds_module_spec", r"M7"),
    ("через сколько оценивают ответ на лечение", "meds_module_spec", r"4–6 неделе"),
]
GOLD_MIN_HITS = 15

OFF_TOPIC = [
    "какой рецепт борща",
    "как починить машину",
    "расписание электричек",
    "кто выиграл чемпионат по футболу",
    "курс доллара сегодня",
    "погода в москве завтра",
    "как установить windows",
    "лучший фильм года",
    "как выучить английский",
]

PATIENT_QUESTIONS = [
    "через сколько начинают действовать антидепрессанты",
    "что такое ПТСР",
    "что такое флэшбек",
    "помогает ли психотерапия при ПТСР",
    "как правильно отменять антидепрессанты",
    "что будет если пить таблетки с алкоголем",
    "сертралин при ПТСР",
    "вызывают ли бензодиазепины зависимость",
]


@pytest.fixture(autouse=True)
def mock_gateway(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(MockTasks()))


class FixedPhraseTasks(MockTasks):
    """An LLM that returns a fixed rephrasing."""

    name = "fake-llm"

    def __init__(self, reply: str) -> None:
        self.reply = reply

    async def phrase(self, intent: str, payload: ScrubbedText, base_text: str) -> str | None:
        return self.reply


def _words(text: str) -> str:
    return " ".join(re.findall(r"\w+", LIT_REFS.sub("", text).lower()))


# --- stemmer ------------------------------------------------------------------------------------------


@pytest.mark.parametrize(
    ("word", "expected"),
    [
        # reference outputs of the Snowball Russian stemmer
        ("вагоне", "вагон"),
        ("важнейшими", "важн"),
        ("гостиница", "гостиниц"),
        ("сделавшись", "сдела"),
        ("посттравматическое", "посттравматическ"),
        ("антидепрессантов", "антидепрессант"),
        ("пробуждения", "пробужден"),
        ("ёлками", "елк"),
        ("SSRI", "ssri"),
    ],
)
def test_stemmer(word: str, expected: str) -> None:
    assert stem(word) == expected


def test_tokens_drop_stop_words_and_split_hyphens() -> None:
    assert tokens("что такое когнитивно-поведенческая терапия") == [
        "когнитивн-поведенческ",
        "когнитивн",
        "поведенческ",
        "терап",
    ]


# --- corpus -------------------------------------------------------------------------------------------


def _page_of(text: str, line: str) -> int | None:
    """Page number of the first line equal to `line`."""
    page = None
    for ln in text.splitlines():
        if m := re.match(r"^--- стр\. (\d+) ---$", ln):
            page = int(m.group(1))
        elif ln == line:
            return page
    return None


def test_every_page_is_chunked_except_bibliography() -> None:
    chunks = load_chunks()
    for doc in load_config().documents:
        text = (rag_dir() / doc.file).read_text(encoding="utf-8")
        n_pages = len(re.findall(r"^--- стр\. \d+ ---$", text, re.M))
        own = [c for c in chunks if c.doc.id == doc.id]
        assert own, doc.id
        assert all(1 <= c.page_start <= c.page_end <= n_pages for c in own)
        missing = set(range(1, n_pages + 1)) - {p for c in own for p in range(c.page_start, c.page_end + 1)}
        lit, appendix = _page_of(text, "Список литературы"), _page_of(text, "Приложение А1")
        allowed = set(range(lit + 1, appendix)) if lit and appendix else set()
        assert missing <= allowed, (doc.id, sorted(missing - allowed))


def test_chunks_are_bounded_and_ids_unique() -> None:
    chunks = load_chunks()
    assert len({c.id for c in chunks}) == len(chunks)
    max_chars = load_config().chunking.max_chars
    # only a single oversized table row may exceed the limit; prose is split by sentences
    for c in chunks:
        if len(c.text) > 2 * max_chars:
            assert any(len(line) > max_chars for line in c.text.splitlines() if line.startswith("|")), c.id


def test_bibliography_and_abbreviations_not_indexed() -> None:
    sections = {c.section for c in load_chunks()}
    assert not any(s.startswith(("Список литературы", "Список сокращений")) for s in sections)
    assert "3.2. Психофармакотерапия" in sections
    assert "3.6 Противорецидивная (профилактическая) терапия" in sections


def test_table_pieces_repeat_header() -> None:
    for c in load_chunks():
        lines = c.text.splitlines()
        for i, line in enumerate(lines):
            if line.startswith("|") and (i == 0 or not lines[i - 1].startswith("|")):
                assert i + 1 < len(lines) and re.match(r"^\|(-+\|)+$", lines[i + 1]), c.id


def test_hamd_items_not_duplicated_or_split() -> None:
    """Converter regression: item 1 of HAM-D was repeated as a fake header on every page, items 2/17 split."""
    text = (rag_dir() / "corpus" / "kr_depression.txt").read_text(encoding="utf-8")
    assert text.count("| 1. СНИЖЕННОЕ НАСТРОЕНИЕ") == 1
    assert re.search(r"^\| 2\. ЧУВСТВО ВИНЫ 0 □", text, re.M)
    assert re.search(r"^\| 17\. КРИТИКА 0 □", text, re.M)


@pytest.mark.parametrize(
    ("line", "is_heading"),
    [
        ("2.1 Жалобы и анамнез", True),
        ("3.1.2. Купирующая терапия ДЭ средней степени", True),
        ("3. Лечение, включая медикаментозную и немедикаментозную терапии", True),
        ("Приложение Г5", True),
        ("2. Стойкие симптомы повышения психологической чувствительности", False),  # list item
        ("28.07.2022 N 69435).", False),  # date
        ("| 2.1 | таблица |", False),
    ],
)
def test_heading(line: str, is_heading: bool) -> None:
    assert (heading(line) is not None) == is_heading


def test_split_sentences_keeps_abbreviations() -> None:
    assert split_sentences("Алгоритм ведения (см. Приложение Б5). Далее текст.") == [
        "Алгоритм ведения (см. Приложение Б5).",
        "Далее текст.",
    ]


# --- retrieval ----------------------------------------------------------------------------------------


def test_gold_questions_hit_at_3() -> None:
    index = get_index()
    misses = [
        q
        for q, doc, pattern in GOLD
        if not any(
            h.chunk.doc.id == doc and re.search(pattern, h.chunk.text) for h in index.search(q, top_k=3, min_score=0)
        )
    ]
    assert len(GOLD) - len(misses) >= GOLD_MIN_HITS, misses


def test_synonyms_expand_abbreviations() -> None:
    terms = get_index().query_terms("ПТСР")
    assert terms["птср"] == 1.0
    assert 0 < terms["посттравматическ"] < 1.0


def test_team_documents_never_reach_patients_or_doctors() -> None:
    for audience in ("patient", "doctor"):
        hits = get_index().search("правила M1–M10 модуль лекарства", audience=audience, min_score=0, top_k=20)
        assert all("team" not in h.chunk.doc.audiences for h in hits)
    team = get_index().search("правило M7 срочные симптомы", audience="team")
    assert team and team[0].chunk.doc.id == "meds_module_spec"


def test_patient_information_appendix_boosted() -> None:
    hit = get_index().search("что такое ПТСР", audience="patient")[0]
    assert hit.chunk.section.startswith("Приложение В")


# --- answers ------------------------------------------------------------------------------------------


@pytest.mark.parametrize("question", OFF_TOPIC)
async def test_off_topic_has_no_answer(question: str) -> None:
    assert await answer(S(question), audience="patient") is None
    assert await answer(S(question), audience="doctor") is None


@pytest.mark.parametrize("question", ["можно увеличить дозу сертралина?", "можно пить полтаблетки?"])
async def test_dose_question_redirected(question: str) -> None:
    assert await answer(S(question), audience="patient") is None
    assert "врач" in fallback_text("patient")


@pytest.mark.parametrize("question", PATIENT_QUESTIONS)
async def test_patient_answer_is_extractive_and_safe(question: str) -> None:
    reply = await answer(S(question), audience="patient")
    if reply is None:  # "ask your doctor" is always a valid outcome
        return
    sources = {c.id: _words(c.text) for c in load_chunks()}
    for sentence in split_sentences(reply.excerpt):
        assert any(_words(sentence) in text for text in sources.values()), sentence  # quoted, not generated
    assert not PRESCRIBING.search(reply.excerpt)
    assert not any(p.search(reply.excerpt) for p in QUOTE_FILTER)
    assert "Источник:" in reply.text and "стр." in reply.text
    assert reply.text.rstrip().endswith(answer_module.load_rag_texts()["patient"].doctor_note)
    assert all(c.doc != "meds_module_spec" for c in reply.citations)


async def test_onset_answer_content() -> None:
    reply = await answer(S("через сколько начинают действовать антидепрессанты"), audience="patient")
    assert reply is not None
    assert re.search(r"недел", reply.excerpt)
    assert reply.coverage >= load_config().answer.min_coverage


async def test_definition_question_prefers_definition() -> None:
    reply = await answer(S("что такое флэшбек"), audience="patient")
    assert reply is not None and reply.excerpt.startswith("Флэшбек (flashback")
    assert reply.citations[0].pages == "2"


async def test_doctor_answer_verbatim_with_citations() -> None:
    reply = await answer(S("синдром отмены антидепрессантов"), audience="doctor")
    assert reply is not None and not reply.phrased_by_llm
    assert re.search(r"\[КР Депрессия 2024, стр\. \d+(–\d+)?\]", reply.text)
    assert "СОПА" in reply.excerpt


async def test_doctor_sees_dosing_patient_does_not() -> None:
    q = S("сертралин при ПТСР")
    doctor = await answer(q, audience="doctor")
    assert doctor is not None and re.search(r"сертралин\*\* 25-200 мг", doctor.excerpt)
    patient = await answer(q, audience="patient")
    assert patient is None or "мг" not in patient.excerpt


async def test_llm_rephrasing_that_adds_numbers_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(FixedPhraseTasks("Эффект будет через 2 дня.")))
    reply = await answer(S("через сколько начинают действовать антидепрессанты"), audience="patient")
    assert reply is not None and not reply.phrased_by_llm
    assert "2 дня" not in reply.text


async def test_llm_rephrasing_kept_when_faithful(monkeypatch: pytest.MonkeyPatch) -> None:
    q = "помогает ли психотерапия при ПТСР"
    excerpt = " ".join(p.text for p in select(get_index(), q, "patient"))
    assert "Это помогает" in excerpt  # anaphora pulled in the sentence it refers to
    simplified = excerpt.replace("Одновременно с этим, психотерапия", "Кроме того, психотерапия")
    monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(FixedPhraseTasks(simplified)))
    reply = await answer(S(q), audience="patient")
    assert reply is not None and reply.phrased_by_llm
    assert simplified in reply.text and reply.excerpt == excerpt  # the verbatim quote is kept for audit