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`.
"""

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] = {}


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.
    """

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

    return deco


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:
    return _triggers.get(key)


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