PTSRaiAgent/app/dialog/flows/meds.py

"""Medication flow (#37, case scenario «управление приёмом лекарств»). Module-agnostic: every patient with a
prescription gets it, whatever the clinical module.

    system:med_reminder {dose_id, drug, dose}
        ──▶ «Принял»                → dose taken, check-in med_taken = true
        ──▶ «Пропущу»               → dose skipped, med_taken = false ──▶ reason:
                                       забыл / побочки (→ side effects) / решил бросить (quit_intent) / другое
        ──▶ «Есть побочные эффекты» → code buttons (тошнота, бессонница, …) or «Другое» → free text
                                       → check-in side_effects (+ dose.side_effects) → «А препарат сегодня приняли?»
    menu «Вопрос о лекарстве" ──▶ meds.question (free text) ──▶ med_qa (instruction excerpt) or «уточните у врача»
    rule action med_info (actions.py hook) ──▶ instruction excerpt for the rule topic + fixed notice

Everything the patient answers is written to the clinical DB (`doses`, `daily_checkins`) and the step asks the
engine to run monitoring (`notify_monitoring`): the rules (lane B, ADR-007) decide alerts and further actions
(quit intent → delayed-onset explanation, missed doses → caring contact). Nothing here decides risk.

Reminder buttons carry the dose they refer to (#56): "med:taken:<ref>", "med:skip:<ref>",
"med:side_effects:<ref>" and the follow-up buttons likewise ("med:se:<code>:<ref>", "med:skip_reason:<r>:<ref>",
"med:dose:yes:<ref>"), matched by prefix triggers. `<ref>` is the `doses.id` of the reminder or, when there is
no dose row (prescription only known to the МИС), the reminder's ISO date. A button of an earlier day's dose is
stale (friendly note, nothing recorded); so are the old payload-less buttons ("med:taken") of earlier releases.
Check-in answers without a button (daily.med step, free text) go to today's planned dose, if any.
"""

from datetime import date, timedelta
from functools import lru_cache
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, ConfigDict
from sqlalchemy import select

from app.config import get_settings
from app.core.privacy import scrub
from app.core.rules import Action
from app.db import repo
from app.db.clinical import DailyCheckin, Dose
from app.db.clinical import Prescription as PrescriptionRow
from app.dialog.actions import register_med_info_handler
from app.dialog.contracts import Button, Event, Reply
from app.dialog.flows.core import menu_reply
from app.dialog.fsm import Ctx, Step, abandonable, get_trigger, handler_for, state, trigger, trigger_payload
from app.integrations.base import integrations
from app.llm import agent
from app.llm.med_qa import DOSE_QUESTION, fallback_text, med_qa

QUESTION_STATE = "meds.question"
SIDE_EFFECT_TEXT_STATE = "meds.side_effect_text"
# A client sends this button when the person submits a free-text step without writing anything.
NO_COMMENT = "text:skip"
SIDE_EFFECT_CODES = ("nausea", "insomnia", "headache", "dizziness", "agitation", "other")
SKIP_REASONS = ("forgot", "side_effects", "quit", "other")


# ------------------------------------------------------------------------------------------- texts


class _Strict(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)


class ReminderTexts(_Strict):
    text: str
    text_no_drug: str
    taken: str
    skip: str
    side_effects: str


class SkipTexts(_Strict):
    question: str
    reasons: dict[str, str]
    ack: dict[str, str]


class SideEffectTexts(_Strict):
    question: str
    options: dict[str, str]
    questions: dict[str, str]
    describe: str
    ack: str
    repeat: str
    ask_dose: str
    dose_yes: str
    dose_no: str


class CheckinMedTexts(_Strict):
    question: str
    question_no_drug: str
    yes: str
    no: str
    later: str


class QuestionTexts(_Strict):
    prompt: str
    no_prescription: str


class MedInfoTopic(_Strict):
    section: str  # med_qa section code (side_effects | onset | alcohol | ...)
    question: str
    lead: str = ""
    fallback: str
    notice: str = ""


class MedsTexts(_Strict):
    status: str
    reminder: ReminderTexts
    taken_ack: str
    stale: str
    skip: SkipTexts
    side_effects: SideEffectTexts
    checkin: CheckinMedTexts
    question: QuestionTexts
    med_info: dict[str, MedInfoTopic]


def meds_texts_path() -> Path:
    return get_settings().config_dir / "texts" / "meds.yaml"


@lru_cache
def load_meds_texts(path: Path | None = None) -> MedsTexts:
    path = path or meds_texts_path()
    return MedsTexts.model_validate(yaml.safe_load(path.read_text(encoding="utf-8")))


# ------------------------------------------------------------------------------------------- data helpers


class Rx(BaseModel):
    """Active prescription as the dialog needs it (drug INN + dose exactly as the doctor wrote them)."""

    id: int | None
    drug: str
    dose: str


async def current_prescription(ctx: Ctx) -> Rx | None:
    """Latest prescription active today: clinical DB first, then the МИС integration (demo: in-memory)."""
    today = ctx.now.date()
    rows = await ctx.session.scalars(
        select(PrescriptionRow).where(PrescriptionRow.token == ctx.token).order_by(PrescriptionRow.id.desc())
    )
    for row in rows:
        if (row.start_date is None or row.start_date <= today) and (row.end_date is None or row.end_date >= today):
            return Rx(id=row.id, drug=row.drug, dose=row.dose)
    for p in await integrations.prescriptions.get_prescriptions(ctx.token):
        if (p.start_date is None or p.start_date <= today) and (p.end_date is None or p.end_date >= today):
            return Rx(id=None, drug=p.drug, dose=p.dose)
    return None


async def checkin_of(ctx: Ctx, day: date) -> DailyCheckin | None:
    return await ctx.session.scalar(
        select(DailyCheckin).where(DailyCheckin.token == ctx.token, DailyCheckin.day == day)
    )


async def save_checkin(ctx: Ctx, **fields: Any) -> DailyCheckin:
    return await repo.upsert_checkin(ctx.session, token=ctx.token, day=ctx.now.date(), now=ctx.now, **fields)


def dose_ref(dose_id: int | None, day: date) -> str:
    """Button payload naming a dose: its row id, or the reminder day when there is no row."""
    return str(dose_id) if dose_id is not None else day.isoformat()


async def todays_dose(ctx: Ctx) -> Dose | None:
    """Latest dose row planned for today (local day of `ctx.now`) — target of answers given without a button."""
    rows = await ctx.session.scalars(select(Dose).where(Dose.token == ctx.token).order_by(Dose.id.desc()).limit(20))
    today = ctx.now.date()
    return next((r for r in rows if r.planned_at.astimezone(ctx.now.tzinfo).date() == today), None)


async def resolve_ref(ctx: Ctx, ref: str | None) -> tuple[bool, Dose | None]:
    """(fresh, dose row) for a button payload. Not fresh: missing/foreign/unknown ref or a dose of another day."""
    if not ref:
        return False, None
    if ref.isdigit():
        row = await ctx.session.get(Dose, int(ref))
        if row is None or row.token != ctx.token:
            return False, None
        return row.planned_at.astimezone(ctx.now.tzinfo).date() == ctx.now.date(), row
    try:
        return date.fromisoformat(ref) == ctx.now.date(), None
    except ValueError:
        return False, None


async def _dose_row(ctx: Ctx, ref: str | None) -> Dose | None:
    if ref is None:
        return await todays_dose(ctx)
    return (await resolve_ref(ctx, ref))[1]


async def record_dose(ctx: Ctx, taken: bool, *, status: str | None = None, ref: str | None = None) -> None:
    """Dose answer → `doses` row (`ref` from the button, else today's) + `daily_checkins.med_taken` (rules read
    both)."""
    row = await _dose_row(ctx, ref)
    if row is not None:
        row.status = status or ("taken" if taken else "skipped")
        row.answered_at = ctx.now
    await save_checkin(ctx, med_taken=taken)


async def record_side_effects(ctx: Ctx, codes: list[str], *, ref: str | None = None) -> None:
    """Merge side-effect codes into today's check-in (rule `side_effect_3d` counts days with any code)."""
    today = await checkin_of(ctx, ctx.now.date())
    merged = list(dict.fromkeys([*(today.side_effects if today and today.side_effects else []), *codes]))
    await save_checkin(ctx, side_effects=merged)
    row = await _dose_row(ctx, ref)
    if row is not None:
        row.side_effects = ",".join(c for c in dict.fromkeys([*(row.side_effects or "").split(","), *codes]) if c)


async def side_effect_reply(ctx: Ctx, codes: list[str]) -> Reply:
    """Reaction to newly reported side effects.

    Same effect as yesterday → short support (the instruction was already shown; persistence over 3 days is the
    rules' business). New effect with a known instruction question → the excerpt of the drug instruction
    (med_qa, section side_effects) — never generated advice. Otherwise a plain acknowledgement.
    """
    t = load_meds_texts().side_effects
    yesterday = await checkin_of(ctx, ctx.now.date() - timedelta(days=1))
    if yesterday is not None and set(codes) & set(yesterday.side_effects or []):
        return Reply(text=t.repeat, kind="info")
    rx = await current_prescription(ctx)
    for code in codes:
        if code in t.questions and rx is not None:
            answer = await med_qa(rx.drug, scrub(t.questions[code]), topic="side_effects")
            if answer is not None:
                return Reply(text=answer.text, kind="info")
    return Reply(text=t.ack, kind="info")


async def leave_to_idle(ctx: Ctx, event: Event, data: dict[str, Any] | None = None) -> Step:
    """Abandon the current sub-flow and let the idle state handle the event (menu button, system event...)."""
    step = Step(next_state="idle", data=data or {})
    idle = handler_for("idle")
    return step.then(await idle(ctx, event)) if idle is not None else step


async def run_trigger_inline(ctx: Ctx, key: str, event: Event) -> Step | None:
    """Run an idle trigger from inside another flow (e.g. a reminder button pressed during the check-in)."""
    handler = get_trigger(key)
    return await handler(ctx, event) if handler is not None else None


# ------------------------------------------------------------------------------------------- reminder


def _reminder_buttons(ref: str) -> list[list[Button]]:
    t = load_meds_texts().reminder
    return [
        [Button(label=t.taken, data=f"med:taken:{ref}")],
        [Button(label=t.skip, data=f"med:skip:{ref}")],
        [Button(label=t.side_effects, data=f"med:side_effects:{ref}")],
    ]


@trigger("system:med_reminder", replace=True)
async def _med_reminder(ctx: Ctx, event: Event) -> Step:
    t = load_meds_texts().reminder
    rx = await current_prescription(ctx)
    drug = str(event.params.get("drug") or (rx.drug if rx else ""))
    dose = str(event.params.get("dose") or (rx.dose if rx else ""))
    try:
        dose_id: int | None = int(event.params["dose_id"])  # type: ignore[call-overload]
    except (KeyError, TypeError, ValueError):
        dose_id = None
    if dose_id is None and rx is not None and rx.id is not None:
        # No scheduler row (demo console without #13): create the planned dose so the answer has a record.
        row = Dose(token=ctx.token, prescription_id=rx.id, planned_at=ctx.now, status="pending")
        ctx.session.add(row)
        await ctx.session.flush()
        dose_id = row.id
    text = t.text.format(drug=drug, dose=dose).replace("  ", " ") if drug else t.text_no_drug
    reply = Reply(text=text, buttons=_reminder_buttons(dose_ref(dose_id, ctx.now.date())), kind="question")
    return Step(replies=[reply], data={"med_pending": None})  # key of earlier releases: the button carries it now


def _stale_reply() -> Step:
    return Step(replies=[Reply(text=load_meds_texts().stale)])


async def _legacy(ctx: Ctx, event: Event) -> Step:
    """Payload-less medication button of an earlier release: the dose is unknown → treat as stale."""
    return _stale_reply()


def _with_dose(fn: Any) -> Any:
    """Prefix-trigger handler wrapper: resolve the dose payload, answer stale buttons without recording."""

    async def handler(ctx: Ctx, event: Event) -> Step:
        ref = trigger_payload(event)
        fresh, _row = await resolve_ref(ctx, ref)
        if not fresh:
            return _stale_reply()
        return await fn(ctx, event, ref)

    return handler


def med_trigger(name: str) -> Any:
    """Register `button:med:<name>:*` (with the dose payload) plus the legacy payload-less key as stale."""

    def deco(fn: Any) -> Any:
        trigger(f"button:med:{name}:*")(_with_dose(fn))
        trigger(f"button:med:{name}")(_legacy)
        return fn

    return deco


@med_trigger("taken")
async def _taken(ctx: Ctx, event: Event, ref: str) -> Step:
    await record_dose(ctx, True, ref=ref)
    return Step(replies=[Reply(text=load_meds_texts().taken_ack)], notify_monitoring=True)


@med_trigger("skip")
async def _skip(ctx: Ctx, event: Event, ref: str) -> Step:
    await record_dose(ctx, False, status="skipped", ref=ref)
    t = load_meds_texts().skip
    buttons = [[Button(label=t.reasons[r], data=f"med:skip_reason:{r}:{ref}")] for r in SKIP_REASONS]
    return Step(replies=[Reply(text=t.question, buttons=buttons, kind="question")], notify_monitoring=True)


def _side_effect_question(ref: str) -> Reply:
    t = load_meds_texts().side_effects
    buttons = [[Button(label=t.options[c], data=f"med:se:{c}:{ref}")] for c in SIDE_EFFECT_CODES]
    return Reply(text=t.question, buttons=buttons, kind="question")


def _skip_reason(reason: str) -> Any:
    async def handler(ctx: Ctx, event: Event, ref: str) -> Step:
        if reason == "side_effects":
            return Step(replies=[_side_effect_question(ref)])
        if reason == "quit":
            # «Решил бросить» is the same signal as "брошу" in free text: rule `quit_intent` reacts (alert +
            # delayed-onset explanation through the med_info action).
            await save_checkin(ctx, quit_intent=True)
            return Step(replies=[Reply(text=load_meds_texts().skip.ack["quit"])], notify_monitoring=True)
        return Step(replies=[Reply(text=load_meds_texts().skip.ack[reason])])

    return handler


for _reason in SKIP_REASONS:
    med_trigger(f"skip_reason:{_reason}")(_skip_reason(_reason))


@med_trigger("side_effects")
async def _side_effects(ctx: Ctx, event: Event, ref: str) -> Step:
    return Step(replies=[_side_effect_question(ref)])


async def _after_side_effects(ctx: Ctx, codes: list[str], ref: str | None) -> Step:
    if ref is None:
        row = await todays_dose(ctx)
        ref = dose_ref(row.id if row is not None else None, ctx.now.date())
    await record_side_effects(ctx, codes, ref=ref)
    replies = [await side_effect_reply(ctx, codes)]
    today = await checkin_of(ctx, ctx.now.date())
    if (today is None or today.med_taken is None) and await current_prescription(ctx) is not None:
        t = load_meds_texts().side_effects
        buttons = [
            [Button(label=t.dose_yes, data=f"med:dose:yes:{ref}")],
            [Button(label=t.dose_no, data=f"med:dose:no:{ref}")],
        ]
        replies.append(Reply(text=t.ask_dose, buttons=buttons, kind="question"))
    return Step(replies=replies, next_state="idle", data={"med_se_ref": None}, notify_monitoring=True)


def _side_effect_code(code: str) -> Any:
    async def handler(ctx: Ctx, event: Event, ref: str) -> Step:
        if code == "other":
            return Step(
                # kind="question": the engine waits for free text here (clients show an input, optional).
                replies=[Reply(text=load_meds_texts().side_effects.describe, kind="question")],
                next_state=SIDE_EFFECT_TEXT_STATE,
                data={"med_se_ref": ref},
            )
        return await _after_side_effects(ctx, [code], ref)

    return handler


for _code in SIDE_EFFECT_CODES:
    med_trigger(f"se:{_code}")(_side_effect_code(_code))


@state(SIDE_EFFECT_TEXT_STATE)
async def _side_effect_text(ctx: Ctx, event: Event) -> Step:
    ref = ctx.data.get("med_se_ref")
    if event.type == "button" and event.button == NO_COMMENT:
        # The description is optional: «other» side effect without words.
        return await _after_side_effects(ctx, ["other"], str(ref) if ref else None)
    if event.type != "text" or ctx.text is None:
        return await leave_to_idle(ctx, event, {"med_se_ref": None})
    facts = await ctx.gateway.extract_checkin(ctx.text)
    # Store only normalised codes, never the description itself (it stays in the scrubbed message log).
    codes = [str(c) for c in facts.side_effects] if facts is not None and facts.side_effects else ["other"]
    return await _after_side_effects(ctx, codes, str(ref) if ref else None)


def _dose_answer(taken: bool) -> Any:
    async def handler(ctx: Ctx, event: Event, ref: str) -> Step:
        await record_dose(ctx, taken, status=None if taken else "missed", ref=ref)
        return Step(
            replies=[Reply(text=load_meds_texts().taken_ack if taken else load_meds_texts().skip.ack["other"])],
            notify_monitoring=True,
        )

    return handler


med_trigger("dose:yes")(_dose_answer(True))
med_trigger("dose:no")(_dose_answer(False))

# Free-text waiting states of this flow: a long-forgotten question must not swallow tomorrow's message.
abandonable("meds")


# ------------------------------------------------------------------------------------------- question menu


@trigger("button:menu:med_question", replace=True)
async def _menu_question(ctx: Ctx, event: Event) -> Step:
    t = load_meds_texts().question
    if await current_prescription(ctx) is None:
        return Step(replies=[menu_reply(ctx, t.no_prescription)])
    return Step(replies=[Reply(text=t.prompt, kind="question")], next_state=QUESTION_STATE)


async def answer_question(ctx: Ctx, rx: Rx | None) -> str:
    """Answer from the instruction excerpt only (SPEC §7 п.5); dose questions and gaps → «уточните у врача»."""
    if rx is None or ctx.text is None:
        return load_meds_texts().question.no_prescription
    if DOSE_QUESTION.search(ctx.text.lower()) is None:
        # Agent: instruction + ЕМИАС record (other drugs, conditions, intolerances) + guidelines; else excerpt only.
        smart = await agent.answer_chat(ctx.token, ctx.text)
        if smart is not None:
            return agent.render(smart)
    answer = await med_qa(rx.drug, ctx.text)
    return answer.text if answer is not None else fallback_text()


@state(QUESTION_STATE)
async def _question(ctx: Ctx, event: Event) -> Step:
    if event.type != "text" or ctx.text is None:
        return await leave_to_idle(ctx, event)
    text = await answer_question(ctx, await current_prescription(ctx))
    return Step(replies=[menu_reply(ctx, text)], next_state="idle")


# ------------------------------------------------------------------------------------------- rule actions


async def med_info_reply(ctx: Ctx, action: Action) -> Reply | None:
    """`med_info` action of a rule (quit_intent → onset, alcohol_3d → alcohol): instruction excerpt framed by
    fixed texts. The rule already decided to say it; this only picks the words. None → engine template."""
    spec = load_meds_texts().med_info.get(str(action.params.get("topic")))
    if spec is None:
        return None
    rx = await current_prescription(ctx)
    answer = await med_qa(rx.drug, scrub(spec.question), topic=spec.section) if rx is not None else None
    body = answer.text if answer is not None else spec.fallback
    return Reply(text="\n\n".join(p for p in (spec.lead, body, spec.notice) if p), kind="info")


register_med_info_handler(med_info_reply)