PTSRaiAgent/app/dialog/fsm.py

"""FSM framework of the dialog engine: context, step result and the state / trigger registries.

States are plain strings, namespaced by module ("veteran.onboarding.pcl5"); core states have short names
("new", "consent", "age_gate", "idle", "crisis", "consent_revoked", ...). A flow module
(`app/dialog/flows/<module>.py`) registers its states at import time — the engine never has to be edited to
add a scenario:

    from app.dialog.fsm import Ctx, Step, on_enter, state
    from app.dialog.scale_session import start_scale

    @on_enter("veteran.onboarding.start")          # runs when the FSM *enters* the state
    async def _start(ctx: Ctx) -> Step:
        return await start_scale(ctx, "pcptsd5", return_state="veteran.onboarding.pcl5")

    @state("veteran.onboarding.trauma")            # handles events while the FSM *is in* the state
    async def _trauma(ctx: Ctx, event: Event) -> Step:
        ...
        return Step(replies=[...], next_state="idle", data={"trauma": True})

Name resolution: an exact registration wins, otherwise the longest registered dotted prefix handles the state
("scale.phq9" → the generic "scale" handler). Triggers (`@trigger("button:menu:checkin")`,
`@trigger("system:daily_checkin")`, `@trigger("text:idle")`) let flows take over what the core `idle` state does
with menu buttons, system events and free text; the core registers defaults, a flow overrides with
`replace=True`.

Prefix triggers carry data in the button (#56): `@trigger("button:med:taken:*")` handles every button whose data
starts with "med:taken:"; the rest of the data ("med:taken:42" → "42") reaches the handler as
`trigger_payload(event)`. An exact key always wins over a prefix, a longer prefix over a shorter one.

Abandoned steps (#56): a flow marks its question states with `abandonable("daily")`; if the patient stays silent
in such a state longer than `settings.abandon_after_hours`, the engine returns the FSM to `idle` before
handling the next event (answers already stored stay stored).
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.types import ScrubbedText
from app.db import repo
from app.db.clinical import PatientProfile
from app.dialog.contracts import Event, Reply
from app.dialog.texts import DialogTexts
from app.llm.gateway import LLMGateway
from app.services.contracts import MonitoringService

# States where nothing is being asked: a pending scale may start here (see engine `_start_pending`).
RESTING_STATES = frozenset({"idle", "crisis"})
# States without (valid) consent: the engine stores nothing and never calls the LLM here (152-FZ art. 9-10).
NO_CONSENT_STATES = frozenset({"new", "consent", "refused", "minor", "consent_revoked"})


@dataclass
class Ctx:
    """Everything a state handler may use. `data` is the persisted `dialog_state.data` (treat as read-only:
    return changes in `Step.data`, the engine applies them)."""

    token: str
    state: str
    data: dict[str, Any]
    now: datetime  # from app.core.clock, taken once per incoming event (ADR-005)
    session: AsyncSession  # clinical DB unit of work of this event
    gateway: LLMGateway
    monitoring: MonitoringService
    texts: DialogTexts
    text: ScrubbedText | None = None  # scrubbed patient text for text events — the only text handlers see
    _profile: PatientProfile | None = field(default=None, repr=False)
    _profile_loaded: bool = field(default=False, repr=False)

    async def profile(self) -> PatientProfile | None:
        if not self._profile_loaded:
            self._profile = await repo.get_profile(self.session, self.token)
            self._profile_loaded = True
        return self._profile


@dataclass
class Step:
    """Result of a handler.

    replies          — messages to send, in order;
    next_state       — new FSM state (None = stay); entering a different state runs its `@on_enter` hook;
    data             — updates merged into `dialog_state.data` (a value of None deletes the key);
    notify_monitoring — the handler wrote clinical data: the engine commits and calls
                       `MonitoringService.after_data_change()`, then executes the returned dialog actions.
    """

    replies: list[Reply] = field(default_factory=list)
    next_state: str | None = None
    data: dict[str, Any] = field(default_factory=dict)
    notify_monitoring: bool = False

    def then(self, other: Step) -> Step:
        """Sequential composition: replies concatenated, later state/data win."""
        return Step(
            replies=[*self.replies, *other.replies],
            next_state=other.next_state or self.next_state,
            data={**self.data, **other.data},
            notify_monitoring=self.notify_monitoring or other.notify_monitoring,
        )


Handler = Callable[[Ctx, Event], Awaitable[Step]]
EnterHook = Callable[[Ctx], Awaitable[Step]]

_handlers: dict[str, Handler] = {}
_enter_hooks: dict[str, EnterHook] = {}
_triggers: dict[str, Handler] = {}
_prefix_triggers: dict[str, Handler] = {}  # "button:med:taken:" → handler (registered as "button:med:taken:*")
_abandonable: dict[str, Callable[[Ctx], bool] | None] = {}

PREFIX_WILDCARD = "*"
PAYLOAD_PARAM = "trigger_payload"


def state(name: str) -> Callable[[Handler], Handler]:
    """Register the event handler of state `name` (or of a dotted prefix, e.g. "scale")."""

    def deco(fn: Handler) -> Handler:
        if name in _handlers:
            raise ValueError(f"state '{name}' already registered")
        _handlers[name] = fn
        return fn

    return deco


def on_enter(name: str) -> Callable[[EnterHook], EnterHook]:
    """Register a hook that runs when the FSM enters state `name` (asks the first question, chains flows)."""

    def deco(fn: EnterHook) -> EnterHook:
        if name in _enter_hooks:
            raise ValueError(f"enter hook of '{name}' already registered")
        _enter_hooks[name] = fn
        return fn

    return deco


def trigger(key: str, *, replace: bool = False) -> Callable[[Handler], Handler]:
    """Register a handler for an idle-menu button / system event / idle free text.

    Keys: "button:<Button.data>", "system:<SystemEventKind>", "text:idle". The core registers defaults;
    flows pass `replace=True` to take one over. A key ending in ":*" is a prefix trigger (see module doc).
    """
    table, name = _triggers, key
    if key.endswith(":" + PREFIX_WILDCARD):
        table, name = _prefix_triggers, key[: -len(PREFIX_WILDCARD)]

    def deco(fn: Handler) -> Handler:
        if name in table and not replace:
            raise ValueError(f"trigger '{key}' already registered (use replace=True to override)")
        table[name] = fn
        return fn

    return deco


def trigger_payload(event: Event) -> str | None:
    """Data after the matched prefix of a prefix trigger ("med:taken:42" → "42"); None for exact triggers."""
    value = event.params.get(PAYLOAD_PARAM)
    return str(value) if value is not None else None


def abandonable(prefix: str, when: Callable[[Ctx], bool] | None = None) -> None:
    """Mark states `prefix` / `prefix.*` as question steps the engine may abandon after a long silence.

    `when(ctx)` may veto it for a particular situation (e.g. a scale that is part of an onboarding).
    """
    if prefix in _abandonable:
        raise ValueError(f"abandonable '{prefix}' already registered")
    _abandonable[prefix] = when


def is_abandonable(ctx: Ctx) -> bool:
    key = _resolve(ctx.state, _abandonable)
    if key is None:
        return False
    when = _abandonable[key]
    return when is None or when(ctx)


def _resolve(name: str, table: dict[str, Any]) -> str | None:
    candidate = name
    while candidate:
        if candidate in table:
            return candidate
        if "." not in candidate:
            return None
        candidate = candidate.rsplit(".", 1)[0]
    return None


def handler_for(name: str) -> Handler | None:
    key = _resolve(name, _handlers)
    return _handlers[key] if key else None


def enter_hook_for(name: str) -> EnterHook | None:
    key = _resolve(name, _enter_hooks)
    return _enter_hooks[key] if key else None


def get_trigger(key: str) -> Handler | None:
    """Exact trigger, else the longest matching prefix trigger (wrapped so it sees the payload)."""
    exact = _triggers.get(key)
    if exact is not None:
        return exact
    matches = [p for p in _prefix_triggers if key.startswith(p) and len(key) > len(p)]
    if not matches:
        return None
    prefix = max(matches, key=len)
    fn, payload = _prefix_triggers[prefix], key[len(prefix) :]

    async def with_payload(ctx: Ctx, event: Event) -> Step:
        return await fn(ctx, event.model_copy(update={"params": {**event.params, PAYLOAD_PARAM: payload}}))

    return with_payload


def is_registered(name: str) -> bool:
    """Exact registration (handler or enter hook) — used to find a module's onboarding entry point."""
    return name in _handlers or name in _enter_hooks


def apply_data(data: dict[str, Any], updates: dict[str, Any]) -> None:
    for key, value in updates.items():
        if value is None:
            data.pop(key, None)
        else:
            data[key] = value