PTSRaiAgent/app/dialog/scale_session.py

"""Reusable scale session: asks a questionnaire item by item inside the dialog (SPEC §6 "Сбор шкалы в диалоге").

Usage from a flow (#36/#37):

    step = await start_scale(ctx, "pcl5", return_state="veteran.onboarding.phq9", scheduled=True)
    return step  # FSM goes to "scale.pcl5"; its enter hook asks item 1; on completion → return_state

Algorithm per item:
  1. The item is shown with one button per answer option from `config/scales/<id>.yaml`.
  2. A button press is taken as is (source="button").
  3. Free text is accepted only for items without `buttons_only` (PHQ-9 item 9 and the whole C-SSRS are
     buttons-only: suicidal-ideation answers are never interpreted by the LLM — SPEC §6, ADR-002). The scrubbed
     text goes to `gateway.extract_item`; `None` (failure / confidence < 0.5 / out-of-range score) → a clarifying
     question with the same buttons, nothing is stored (source="llm" only for a confident extraction).
  4. Every answer is stored immediately (`item_responses`), so a session interrupted by a crisis keeps its data.
  5. After each answer the scale is scored deterministically (`app.core.scales.score`). If a crisis flag from
     `config/texts/dialog.yaml: settings.crisis_flags` fires (PHQ-9 item 9 ≥ 1, C-SSRS high) the session ends at
     once: partial result saved, fixed crisis script, C-SSRS queued (unless this *is* the C-SSRS, or a C-SSRS was
     completed less than `settings.cssrs_cooldown_hours` ago — the patient is not asked the same suicide screen
     twice in a day; the crisis script and the alert still happen). A positive item 9 is acted upon even in a
     half-finished questionnaire.
  6. On completion: `scale_results` row (total, severity, item 9, scheduled flag, flags) → the engine commits
     and calls `MonitoringService.after_data_change()` (rules decide alerts — ADR-007) → back to `return_state`.
The patient never sees a score or a severity label: that would read as a diagnosis.

Intros (#56): `start_scale(intro=None)` looks the intro up in a registry filled by flows
(`register_scale_intro("phq9", fn_or_text)`), so scales started by the engine (`scale_due`, rule actions) get
the flow's friendly intro without the flow hooking `scale.<id>` itself.
"""

from collections.abc import Callable
from dataclasses import replace
from datetime import timedelta
from typing import Any

from app.core.privacy import scrub
from app.core.scales import Item, ScaleDef, ScaleResult, load_scale, score
from app.db import repo
from app.dialog.contracts import Button, Event, Reply
from app.dialog.crisis import crisis_open, open_crisis
from app.dialog.fsm import RESTING_STATES, Ctx, Step, abandonable, is_abandonable, on_enter, state

PREFIX = "scale"

# (ctx, scale_id, scheduled, return_state) → intro text or None (no intro in this situation).
IntroFn = Callable[[Ctx, str, bool, str], str | None]
_intros: dict[str, str | IntroFn] = {}


def register_scale_intro(scale_id: str, intro: str | IntroFn, *, replace: bool = False) -> None:
    """Register the intro shown before `scale_id` when a session is started without an explicit `intro`."""
    if scale_id in _intros and not replace:
        raise ValueError(f"intro of scale '{scale_id}' already registered (use replace=True to override)")
    _intros[scale_id] = intro


def scale_intro(ctx: Ctx, scale_id: str, *, scheduled: bool, return_state: str) -> str | None:
    intro = _intros.get(scale_id)
    if intro is None or isinstance(intro, str):
        return intro or None
    return intro(ctx, scale_id, scheduled, return_state) or None


def scale_state(scale_id: str) -> str:
    return f"{PREFIX}.{scale_id}"


def active_scale(state_name: str) -> str | None:
    return state_name.split(".", 1)[1] if state_name.startswith(PREFIX + ".") else None


def add_pending(data: dict[str, Any], scale_id: str, scheduled: bool) -> list[dict[str, Any]]:
    """Queue a scale to start at the next resting state (idle/crisis); duplicates are ignored."""
    pending = list(data.get("pending_scales") or [])
    if scale_id not in {p["scale"] for p in pending}:
        pending.append({"scale": scale_id, "scheduled": scheduled})
    return pending


async def start_scale(
    ctx: Ctx, scale_id: str, *, return_state: str, scheduled: bool = True, intro: str | None = None
) -> Step:
    """Begin a session of `scale_id`; returns a Step moving the FSM to `scale.<id>`.

    `scheduled=False` marks an unscheduled administration (triggered by a rule or a safety signal).
    `intro` — optional extra text shown before the scale instruction; None → the registered intro, if any.
    """
    scale = load_scale(scale_id)  # fail fast on an unknown id
    if intro is None:
        intro = scale_intro(ctx, scale_id, scheduled=scheduled, return_state=return_state)
    sessions = dict(ctx.data.get("scales") or {})
    sessions[scale_id] = {
        "idx": 0,
        "answers": {},
        "return_state": return_state,
        "scheduled": scheduled,
        "intro": intro,
        "intro_shown": False,
    }
    target = scale_state(scale_id)
    if ctx.state == target:  # restart in place: the enter hook would not run
        sessions[scale_id]["intro_shown"] = True
        return Step(replies=[_question(ctx, scale, sessions[scale_id], intro=True)], data={"scales": sessions})
    return Step(next_state=target, data={"scales": sessions})


def _question(ctx: Ctx, scale: ScaleDef, session: dict[str, Any], *, intro: bool = False, note: str = "") -> Reply:
    idx = session["idx"]
    item = scale.items[idx]
    parts: list[str] = []
    if note:
        parts.append(note)
    if intro:
        if session.get("intro"):
            parts.append(session["intro"])
        parts.append(ctx.texts.scale.intro.format(count=len(scale.items), instruction=scale.instruction or "").strip())
    parts.append(ctx.texts.scale.question.format(n=idx + 1, total=len(scale.items), text=item.text))
    buttons = [[Button(label=o.label, data=f"{PREFIX}:{scale.id}:{item.id}:{o.value}")] for o in item.options]
    return Reply(text="\n\n".join(parts), buttons=buttons, kind="question")


def _session(ctx: Ctx) -> tuple[ScaleDef, dict[str, Any]] | None:
    scale_id = active_scale(ctx.state)
    session = (ctx.data.get("scales") or {}).get(scale_id) if scale_id else None
    if scale_id is None or session is None:
        return None
    return load_scale(scale_id), session


def _abandon_allowed(ctx: Ctx) -> bool:
    """A silent scale session may be abandoned only if what it returns to may be (idle, a check-in step) —
    never an onboarding chain, which would otherwise be skipped."""
    found = _session(ctx)
    if found is None:
        return True
    back = str(found[1].get("return_state") or "idle")
    if back in RESTING_STATES:
        return True
    return back != ctx.state and is_abandonable(replace(ctx, state=back))


abandonable(PREFIX, _abandon_allowed)


async def recent_risk_screen(ctx: Ctx, screen: str) -> bool:
    """A complete risk screen (C-SSRS) less than `cssrs_cooldown_hours` ago (scale_results, any trigger)."""
    since = ctx.now - timedelta(hours=ctx.texts.settings.cssrs_cooldown_hours)
    results = await repo.list_scale_results(ctx.session, ctx.token, screen)
    return any(r.complete and r.ts >= since for r in results)


@on_enter(PREFIX)
async def _enter(ctx: Ctx) -> Step:
    found = _session(ctx)
    if found is None:
        return Step(next_state="idle")
    scale, session = found
    show_intro = session["idx"] == 0 and not session.get("intro_shown")
    updated = {**session, "intro_shown": True}
    sessions = {**(ctx.data.get("scales") or {}), scale.id: updated}
    return Step(replies=[_question(ctx, scale, updated, intro=show_intro)], data={"scales": sessions})


def _parse_button(data: str, scale: ScaleDef, item: Item) -> int | None:
    parts = data.split(":")
    if len(parts) != 4 or parts[0] != PREFIX or parts[1] != scale.id or parts[2] != item.id:
        return None  # stale button from an earlier question
    try:
        value = int(parts[3])
    except ValueError:
        return None
    return value if value in {o.value for o in item.options} else None


@state(PREFIX)
async def _handle(ctx: Ctx, event: Event) -> Step:
    found = _session(ctx)
    if found is None:
        return Step(next_state="idle")
    scale, session = found
    item = scale.items[session["idx"]]

    if event.type == "button" and event.button:
        value = _parse_button(event.button, scale, item)
        if value is not None:
            return await _record(ctx, scale, session, item, value, source="button")
        return Step(replies=[_question(ctx, scale, session)])

    if event.type == "text" and ctx.text is not None:
        if item.buttons_only:
            # Safety items are never LLM-interpreted: ask again for a button.
            return Step(replies=[_question(ctx, scale, session, note=ctx.texts.scale.buttons_only)])
        extracted = await ctx.gateway.extract_item(item.id, item.text, item.max_value, ctx.text)
        if extracted is None or extracted.score not in {o.value for o in item.options}:
            return Step(replies=[_question(ctx, scale, session, note=ctx.texts.scale.clarify)])
        return await _record(
            ctx,
            scale,
            session,
            item,
            extracted.score,
            source="llm",
            confidence=extracted.confidence,
            evidence=extracted.evidence,
        )

    if event.type == "system":
        return Step()  # reminders etc. do not interrupt a questionnaire
    return Step(replies=[_question(ctx, scale, session)])


async def _record(
    ctx: Ctx,
    scale: ScaleDef,
    session: dict[str, Any],
    item: Item,
    value: int,
    *,
    source: str,
    confidence: float | None = None,
    evidence: str = "",
) -> Step:
    await repo.add_item_response(
        ctx.session,
        token=ctx.token,
        scale=scale.id,
        item_id=item.id,
        score=value,
        source=source,
        ts=ctx.now,
        confidence=confidence,
        # Evidence is a quote of already scrubbed text; scrub again so nothing unscrubbed can be stored.
        evidence=scrub(evidence) if evidence else None,
    )
    answers = {**session["answers"], item.id: value}
    profile = await ctx.profile()
    sex = profile.sex if profile is not None and profile.sex in ("m", "f") else None
    result = score(scale.id, answers, sex=sex)  # type: ignore[arg-type]
    crisis_flags = [f for f in result.flags if f in ctx.texts.settings.crisis_flags]
    if crisis_flags or session["idx"] + 1 >= len(scale.items):
        return await _finish(ctx, scale, session, answers, result, crisis_flags)
    updated = {**session, "answers": answers, "idx": session["idx"] + 1}
    sessions = {**(ctx.data.get("scales") or {}), scale.id: updated}
    return Step(replies=[_question(ctx, scale, updated)], data={"scales": sessions})


async def _finish(
    ctx: Ctx,
    scale: ScaleDef,
    session: dict[str, Any],
    answers: dict[str, int],
    result: ScaleResult,
    crisis_flags: list[str],
) -> Step:
    await repo.add_scale_result(
        ctx.session,
        token=ctx.token,
        scale=scale.id,
        total=result.total,
        severity=result.severity_code,
        complete=result.complete,
        ts=ctx.now,
        # PHQ-9 item 9 is stored separately: any value >= 1 escalates regardless of the total.
        item9=answers.get("phq9_9") if scale.id == "phq9" else None,
        scheduled=bool(session.get("scheduled", True)),
        flags={f: True for f in result.flags},
    )
    sessions = {k: v for k, v in (ctx.data.get("scales") or {}).items() if k != scale.id}
    done = Step(
        replies=[Reply(text=ctx.texts.scale.done)],
        next_state=session["return_state"],
        data={"scales": sessions},
        notify_monitoring=True,
    )
    if not crisis_flags:
        return done

    # Crisis path decided by the scale's own flags (deterministic), not by the LLM.
    risk_screen = ctx.texts.settings.risk_screen_scale
    if crisis_open(ctx.data):
        crisis = Step(next_state="crisis")
    else:
        crisis = open_crisis(ctx, categories=[], source=f"scale:{scale.id}")
    step = done.then(crisis)
    # Same-day repeat guard: a C-SSRS completed within the cooldown already answered the screening question.
    if scale.id != risk_screen and not await recent_risk_screen(ctx, risk_screen):
        step.data["pending_scales"] = add_pending(ctx.data, risk_screen, scheduled=False)
    return step