PTSRaiAgent/tests/test_dialog_engine_gaps.py

"""Engine gaps found by the 13-day run (#56): no same-day C-SSRS repeat, scale intro registry, prefix triggers
with a button payload (medication buttons carry the dose), abandoned-step timeout."""

from collections.abc import AsyncIterator
from datetime import datetime, timedelta
from typing import Any

import pytest
from sqlalchemy import select

from app.config import Settings
from app.core import clock
from app.core.rules import Action
from app.core.scales import load_scale
from app.db import repo, session
from app.db.bootstrap import init_db
from app.db.clinical import Dose, ItemResponse, Prescription
from app.dialog import engine as engine_module
from app.dialog import fsm, scale_session
from app.dialog.contracts import Event, Reply
from app.dialog.engine import STEP_AT, Engine
from app.dialog.flows.meds import load_meds_texts
from app.dialog.flows.veteran_daily import load_daily_texts
from app.dialog.fsm import Ctx, Step, get_trigger, is_registered, trigger, trigger_payload
from app.dialog.texts import load_dialog_texts
from app.llm.gateway import LLMGateway
from app.llm.mock import MockTasks

TOKEN = "g" * 32
OTHER = "h" * 32
MEM = "sqlite+aiosqlite:///:memory:"
NOW = datetime(2026, 10, 7, 9, 0)
TEXTS = load_dialog_texts()
MEDS = load_meds_texts()
DAILY = load_daily_texts()
SCREEN = TEXTS.settings.risk_screen_scale


class ScriptedMonitoring:
    def __init__(self) -> None:
        self.actions: list[Action] = []

    async def after_data_change(self, token: str) -> list[Action]:
        out, self.actions = self.actions, []
        return out

    async def is_in_crisis(self, token: str) -> bool:
        return True


@pytest.fixture
async def mon(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[ScriptedMonitoring]:
    monkeypatch.setattr(clock, "get_settings", lambda: Settings(_env_file=None, demo_clock=True))
    clock.set_now(NOW)
    await session.configure(identity_url=MEM, clinical_url=MEM)
    await init_db()
    async with session.get_clinical_session() as s:
        await repo.upsert_profile(
            s, token=TOKEN, doctor_id="doc", modules=["veteran"], sex="m", birth_year=1990, therapy_start=NOW.date()
        )
        s.add(Prescription(token=TOKEN, drug="сертралин", dose="50 мг", schedule={"times": ["08:00"]}))
        await repo.set_dialog_state(s, token=TOKEN, state="idle", data={"consent": True}, now=clock.now())
    yield ScriptedMonitoring()
    clock.reset()
    await session.configure()


@pytest.fixture
def engine(mon: ScriptedMonitoring) -> Engine:
    return Engine(gateway=LLMGateway(MockTasks()), monitoring=mon)


async def press(e: Engine, data: str) -> list[Reply]:
    return await e.handle(TOKEN, Event(type="button", button=data))


async def text(e: Engine, t: str) -> list[Reply]:
    return await e.handle(TOKEN, Event(type="text", text=t))


async def system(e: Engine, kind: str, **params: Any) -> list[Reply]:
    return await e.handle(TOKEN, Event(type="system", system=kind, params=params))  # type: ignore[arg-type]


async def answer(e: Engine, scale_id: str, values: list[int]) -> list[Reply]:
    out: list[Reply] = []
    for item, value in zip(load_scale(scale_id).items, values, strict=False):
        out += await press(e, f"scale:{scale_id}:{item.id}:{value}")
    return out


def buttons(replies: list[Reply]) -> list[str]:
    return [b.data for r in replies for row in r.buttons for b in row]


async def get_state() -> tuple[str, dict[str, Any]]:
    async with session.get_clinical_session() as s:
        row = await repo.get_dialog_state(s, TOKEN)
        assert row is not None
        return row.state, dict(row.data)


async def results(scale: str) -> list[Any]:
    async with session.get_clinical_session() as s:
        return list(await repo.list_scale_results(s, TOKEN, scale))


async def checkin() -> Any:
    async with session.get_clinical_session() as s:
        rows = await repo.list_checkins(s, TOKEN)
        return rows[-1] if rows else None


async def doses() -> list[Dose]:
    async with session.get_clinical_session() as s:
        return list((await s.scalars(select(Dose).where(Dose.token == TOKEN).order_by(Dose.id))).all())


# ----------------------------------------------------------------------------------------------- 1. C-SSRS repeat


async def test_item9_right_after_cssrs_does_not_repeat_screen(engine: Engine, mon: ScriptedMonitoring) -> None:
    await system(engine, "scale_due", scale=SCREEN, scheduled=False)
    assert (await get_state())[0] == f"scale.{SCREEN}"
    await answer(engine, SCREEN, [0] * 6)
    clock.advance(hours=1)
    await system(engine, "scale_due", scale="phq9", scheduled=False)
    # The rule phq9_item9 also asks for the screen — neither path may repeat it within the cooldown.
    mon.actions = [Action(type="schedule_scale", code="phq9_item9", params={"scale": SCREEN, "scheduled": False})]
    r = await answer(engine, "phq9", [1] * 8 + [1])
    assert any(x.kind == "crisis" for x in r)  # the fixed crisis script still goes out
    state, data = await get_state()
    assert state == "crisis" and data["crisis"]["open"] is True and not data.get("pending_scales")
    assert len(await results(SCREEN)) == 1


async def test_item9_after_old_cssrs_asks_screen_again(engine: Engine) -> None:
    await system(engine, "scale_due", scale=SCREEN, scheduled=False)
    await answer(engine, SCREEN, [0] * 6)
    clock.advance(hours=TEXTS.settings.cssrs_cooldown_hours + 1)
    await system(engine, "scale_due", scale="phq9", scheduled=False)
    r = await answer(engine, "phq9", [1] * 8 + [1])
    assert any(x.kind == "crisis" for x in r)
    assert (await get_state())[0] == f"scale.{SCREEN}" and TEXTS.scale.risk_screen_intro in r[-1].text


async def test_unfinished_cssrs_does_not_count(engine: Engine) -> None:
    await system(engine, "scale_due", scale=SCREEN, scheduled=False)
    await answer(engine, SCREEN, [0, 0])  # left halfway: no complete result
    await system(engine, "scale_due", scale="phq9", scheduled=False)  # queued behind the screen
    async with session.get_clinical_session() as s:
        assert not await scale_session.recent_risk_screen(_ctx(s), SCREEN)


def _ctx(s: Any, state: str = "idle", data: dict[str, Any] | None = None) -> Ctx:
    return Ctx(
        token=TOKEN,
        state=state,
        data=data or {},
        now=clock.now(),
        session=s,
        gateway=LLMGateway(MockTasks()),
        monitoring=ScriptedMonitoring(),
        texts=TEXTS,
    )


# ----------------------------------------------------------------------------------------------- 2. intros


async def test_registered_intro_text_and_callable(engine: Engine, monkeypatch: pytest.MonkeyPatch) -> None:
    seen: list[tuple[str, bool, str]] = []

    def intro(ctx: Ctx, scale_id: str, scheduled: bool, return_state: str) -> str | None:
        seen.append((scale_id, scheduled, return_state))
        return "Вступление к GAD-7"

    monkeypatch.setattr(scale_session, "_intros", dict(scale_session._intros))
    scale_session.register_scale_intro("gad7", intro)
    with pytest.raises(ValueError):
        scale_session.register_scale_intro("gad7", "второе")
    r = await system(engine, "scale_due", scale="gad7", scheduled=True)
    assert r[0].text.startswith("Вступление к GAD-7") and seen == [("gad7", True, "idle")]

    scale_session.register_scale_intro("gad7", "Текстовое вступление", replace=True)
    async with session.get_clinical_session() as s:
        step = await scale_session.start_scale(_ctx(s), "gad7", return_state="idle")
        assert step.data["scales"]["gad7"]["intro"] == "Текстовое вступление"
        explicit = await scale_session.start_scale(_ctx(s), "gad7", return_state="idle", intro="Своё")
        assert explicit.data["scales"]["gad7"]["intro"] == "Своё"


def test_daily_intros_use_registry_not_enter_hooks() -> None:
    for scale_id in DAILY.settings.intro_scales:
        assert not is_registered(f"scale.{scale_id}")  # no per-scale enter hooks any more
        assert scale_id in scale_session._intros


async def test_unscheduled_intro_from_registry(engine: Engine) -> None:
    r = await system(engine, "scale_due", scale="phq9", scheduled=False)
    assert DAILY.scale_intros.unscheduled["phq9"] in r[0].text


# ----------------------------------------------------------------------------------------------- 3. payload buttons


async def test_prefix_trigger_payload_and_precedence(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(fsm, "_triggers", dict(fsm._triggers))
    monkeypatch.setattr(fsm, "_prefix_triggers", dict(fsm._prefix_triggers))

    def handler(tag: str) -> Any:
        async def fn(ctx: Ctx, event: Event) -> Step:
            return Step(data={"tag": tag, "payload": trigger_payload(event)})

        return fn

    trigger("button:zz:a:*")(handler("short"))
    trigger("button:zz:a:b:*")(handler("long"))
    trigger("button:zz:a:exact")(handler("exact"))
    with pytest.raises(ValueError):
        trigger("button:zz:a:*")(handler("dup"))

    async def run(key: str) -> dict[str, Any] | None:
        fn = get_trigger(key)
        if fn is None:
            return None
        return (await fn(None, Event(type="button", button=key.split(":", 1)[1]))).data  # type: ignore[arg-type]

    assert await run("button:zz:a:42") == {"tag": "short", "payload": "42"}
    assert await run("button:zz:a:b:7:x") == {"tag": "long", "payload": "7:x"}
    assert await run("button:zz:a:exact") == {"tag": "exact", "payload": None}
    assert await run("button:zz:a:") is None and await run("button:zz:a") is None


async def test_reminder_buttons_carry_dose_and_no_state_key(engine: Engine) -> None:
    async with session.get_clinical_session() as s:
        await repo.set_dialog_state(
            s, token=TOKEN, state="idle", data={"consent": True, "med_pending": {"dose_id": 1}}, now=clock.now()
        )
    r = await system(engine, "med_reminder")
    [dose] = await doses()
    assert buttons(r) == [f"med:taken:{dose.id}", f"med:skip:{dose.id}", f"med:side_effects:{dose.id}"]
    assert "med_pending" not in (await get_state())[1]
    skip = await press(engine, f"med:skip:{dose.id}")
    assert f"med:skip_reason:side_effects:{dose.id}" in buttons(skip)
    se = await press(engine, f"med:skip_reason:side_effects:{dose.id}")
    await press(engine, f"med:se:nausea:{dose.id}")
    [dose] = await doses()
    assert dose.status == "skipped" and dose.side_effects == "nausea" and "med:se:nausea:" in buttons(se)[0]


async def test_two_reminders_each_button_answers_its_own_dose(engine: Engine) -> None:
    first = await system(engine, "med_reminder")
    second = await system(engine, "med_reminder")
    await press(engine, buttons(first)[0])  # «Принял» of the FIRST reminder, pressed after the second arrived
    d1, d2 = await doses()
    assert (d1.status, d2.status) == ("taken", "pending")
    await press(engine, buttons(second)[1])
    assert [d.status for d in await doses()] == ["taken", "skipped"]


@pytest.mark.parametrize("data", ["med:taken", "med:skip", "med:se:nausea", "med:dose:yes", "med:skip_reason:quit"])
async def test_legacy_payloadless_buttons_are_stale(engine: Engine, data: str) -> None:
    await system(engine, "med_reminder")
    r = await press(engine, data)
    assert r[0].text == MEDS.stale and await checkin() is None
    assert [d.status for d in await doses()] == ["pending"]


async def test_foreign_unknown_and_old_dose_buttons_are_stale(engine: Engine) -> None:
    async with session.get_clinical_session() as s:
        rx = Prescription(token=OTHER, drug="сертралин", dose="50 мг", schedule={"times": ["08:00"]})
        s.add(rx)
        await s.flush()
        foreign = Dose(token=OTHER, prescription_id=rx.id, planned_at=clock.now(), status="pending")
        s.add(foreign)
        await s.flush()
        foreign_id = foreign.id
    for data in (f"med:taken:{foreign_id}", "med:taken:9999", "med:taken:garbage", "med:taken:2026-10-06"):
        assert (await press(engine, data))[0].text == MEDS.stale
    old = await system(engine, "med_reminder")
    clock.advance(days=1)
    assert (await press(engine, buttons(old)[0]))[0].text == MEDS.stale
    assert await checkin() is None


async def test_reminder_without_dose_row_uses_day_payload(engine: Engine) -> None:
    async with session.get_clinical_session() as s:
        for row in await s.scalars(select(Prescription)):
            await s.delete(row)
    r = await system(engine, "med_reminder")
    assert buttons(r)[0] == f"med:taken:{NOW.date().isoformat()}"
    r = await press(engine, buttons(r)[0])
    assert r[0].text == MEDS.taken_ack and (await checkin()).med_taken is True and await doses() == []


async def test_checkin_med_step_answers_todays_dose(engine: Engine) -> None:
    await system(engine, "med_reminder")
    await system(engine, "daily_checkin")
    await press(engine, "daily:mood:5")
    await press(engine, "daily:med:yes")
    assert [d.status for d in await doses()] == ["taken"]


# ----------------------------------------------------------------------------------------------- 4. abandon timeout


async def test_abandoned_checkin_returns_to_idle(engine: Engine) -> None:
    await system(engine, "daily_checkin")
    await press(engine, "daily:mood:3")
    assert (await get_state())[0].startswith("daily.")
    clock.advance(hours=TEXTS.settings.abandon_after_hours + 1)
    r = await system(engine, "med_reminder")
    state, data = await get_state()
    assert state == "idle" and "checkin" not in data
    assert len(r) == 1 and r[0].text.startswith(MEDS.reminder.text.split("{")[0])  # no check-in question repeated
    assert (await checkin()).mood == 3  # stored answer survives


async def test_short_pause_keeps_checkin(engine: Engine) -> None:
    await system(engine, "daily_checkin")
    clock.advance(hours=TEXTS.settings.abandon_after_hours - 1)
    r = await press(engine, "daily:mood:6")
    assert (await get_state())[0] != "idle" and r


async def test_abandoned_scale_lets_queued_scale_start(engine: Engine) -> None:
    await system(engine, "scale_due", scale="phq9", scheduled=True)
    await answer(engine, "phq9", [1, 2])
    clock.advance(hours=4)
    await system(engine, "scale_due", scale="pcl5", scheduled=True)  # queued; system events don't extend the step
    clock.advance(hours=3)
    r = await text(engine, "добрый вечер")
    state, data = await get_state()
    assert state == "scale.pcl5" and "phq9" not in data["scales"] and not data.get("pending_scales")
    assert DAILY.scale_intros.scheduled["pcl5"] in r[-1].text
    assert await results("phq9") == []  # unfinished: no total, no severity
    async with session.get_clinical_session() as s:
        stored = (await s.scalars(select(ItemResponse).where(ItemResponse.scale == "phq9"))).all()
    assert [x.score for x in stored] == [1, 2]


async def test_timeout_never_applies_in_crisis(engine: Engine) -> None:
    await system(engine, "scale_due", scale="phq9", scheduled=False)
    await answer(engine, "phq9", [0] * 8 + [2])  # item 9 → crisis + C-SSRS (none done before)
    state, data = await get_state()
    assert state == f"scale.{SCREEN}" and data["crisis"]["open"] is True
    clock.advance(hours=TEXTS.settings.abandon_after_hours * 3)
    await press(engine, f"scale:{SCREEN}:cssrs_1:0")
    state, data = await get_state()
    assert state == f"scale.{SCREEN}" and data["scales"][SCREEN]["idx"] == 1


async def test_timeout_never_skips_onboarding_scale(engine: Engine) -> None:
    item = load_scale("pcl5").items[0]
    old = (clock.now() - timedelta(hours=48)).isoformat()
    session_data = {
        "idx": 0,
        "answers": {},
        "return_state": "veteran.onboarding.phq9",
        "scheduled": True,
        "intro": None,
        "intro_shown": True,
    }
    async with session.get_clinical_session() as s:
        await repo.set_dialog_state(
            s,
            token=TOKEN,
            state="scale.pcl5",
            data={"consent": True, "scales": {"pcl5": session_data}, STEP_AT: old},
            now=clock.now(),
        )
    await press(engine, f"scale:pcl5:{item.id}:1")
    state, data = await get_state()
    assert state == "scale.pcl5" and data["scales"]["pcl5"]["idx"] == 1


async def test_timeout_disabled_same_day(engine: Engine, monkeypatch: pytest.MonkeyPatch) -> None:
    clock.set_now(NOW.replace(hour=1))
    off = TEXTS.model_copy(update={"settings": TEXTS.settings.model_copy(update={"abandon_after_hours": 0})})
    monkeypatch.setattr(engine_module, "load_dialog_texts", lambda: off)
    await system(engine, "daily_checkin")
    clock.advance(hours=20)
    await press(engine, "daily:mood:4")
    assert (await get_state())[0] != "idle" and (await checkin()).mood == 4