PTSRaiAgent/app/db/repo.py

"""Thin clinical-DB repository helpers. Must not import `app.db.identity` (guardrail #5).

Functions work inside the caller's session (`get_clinical_session()`), flush but never commit, and take
time as an explicit argument (ADR-005).
"""

from collections.abc import Sequence
from datetime import date, datetime
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.types import AlertLevel, ScrubbedText
from app.db.audit import audit
from app.db.clinical import (
    Alert,
    ClockState,
    DailyCheckin,
    DialogState,
    ItemResponse,
    Message,
    PatientProfile,
    ScaleResult,
)

_CHECKIN_FIELDS = frozenset(
    {
        "mood",
        "med_taken",
        "sleep_bucket",
        "alcohol_bucket",
        "side_effects",
        "quit_intent",
        "nightmares",
        "free_text_scrubbed",
    }
)
_PROFILE_FIELDS = frozenset(
    {
        "doctor_id",
        "modules",
        "therapy_start",
        "next_visit",
        "birth_year",
        "sex",
        "enhanced_monitoring",
        "status_color",
    }
)


# --- patient profiles -------------------------------------------------------------------------------


async def get_profile(session: AsyncSession, token: str) -> PatientProfile | None:
    return await session.get(PatientProfile, token)


async def upsert_profile(session: AsyncSession, *, token: str, **fields: Any) -> PatientProfile:
    """Create or update a profile; only the passed fields are changed."""
    unknown = set(fields) - _PROFILE_FIELDS
    if unknown:
        raise ValueError(f"unknown profile fields: {sorted(unknown)}")
    profile = await session.get(PatientProfile, token)
    if profile is None:
        profile = PatientProfile(token=token, modules=[], enhanced_monitoring=False)
        session.add(profile)
    for key, value in fields.items():
        setattr(profile, key, value)
    await session.flush()
    return profile


async def patients_of_doctor(session: AsyncSession, doctor_id: str) -> Sequence[PatientProfile]:
    rows = await session.scalars(
        select(PatientProfile).where(PatientProfile.doctor_id == doctor_id).order_by(PatientProfile.token)
    )
    return rows.all()


# --- daily check-ins and scales ---------------------------------------------------------------------


async def upsert_checkin(session: AsyncSession, *, token: str, day: date, now: datetime, **fields: Any) -> DailyCheckin:
    """Merge answers into the (token, day) check-in.

    The daily dialog fills the row step by step (mood, then medication, then the module question), so
    only fields passed and not None overwrite stored values.
    """
    unknown = set(fields) - _CHECKIN_FIELDS
    if unknown:
        raise ValueError(f"unknown check-in fields: {sorted(unknown)}")
    row = await session.scalar(select(DailyCheckin).where(DailyCheckin.token == token, DailyCheckin.day == day))
    if row is None:
        row = DailyCheckin(token=token, day=day, side_effects=[], updated_at=now)
        session.add(row)
    for key, value in fields.items():
        if value is not None:
            setattr(row, key, value)
    row.updated_at = now
    await session.flush()
    return row


async def list_checkins(session: AsyncSession, token: str, since: date | None = None) -> Sequence[DailyCheckin]:
    stmt = select(DailyCheckin).where(DailyCheckin.token == token)
    if since is not None:
        stmt = stmt.where(DailyCheckin.day >= since)
    return (await session.scalars(stmt.order_by(DailyCheckin.day))).all()


async def add_item_response(
    session: AsyncSession,
    *,
    token: str,
    scale: str,
    item_id: str,
    score: int | None,
    source: str,
    ts: datetime,
    confidence: float | None = None,
    evidence: ScrubbedText | None = None,
) -> ItemResponse:
    row = ItemResponse(
        token=token,
        scale=scale,
        item_id=item_id,
        score=score,
        source=source,
        ts=ts,
        confidence=confidence,
        evidence=evidence,
    )
    session.add(row)
    await session.flush()
    return row


async def add_scale_result(
    session: AsyncSession,
    *,
    token: str,
    scale: str,
    total: int,
    severity: str,
    complete: bool,
    ts: datetime,
    item9: int | None = None,
    scheduled: bool = True,
    flags: dict[str, Any] | None = None,
    period_start: date | None = None,
    period_end: date | None = None,
) -> ScaleResult:
    row = ScaleResult(
        token=token,
        scale=scale,
        total=total,
        severity=severity,
        complete=complete,
        ts=ts,
        item9=item9,
        scheduled=scheduled,
        flags=flags or {},
        period_start=period_start,
        period_end=period_end,
    )
    session.add(row)
    await session.flush()
    return row


async def list_scale_results(session: AsyncSession, token: str, scale: str | None = None) -> Sequence[ScaleResult]:
    stmt = select(ScaleResult).where(ScaleResult.token == token)
    if scale is not None:
        stmt = stmt.where(ScaleResult.scale == scale)
    return (await session.scalars(stmt.order_by(ScaleResult.ts))).all()


# --- alerts -----------------------------------------------------------------------------------------


async def add_alert(
    session: AsyncSession,
    *,
    token: str,
    level: AlertLevel,
    code: str,
    reason: str,
    ts: datetime,
    params: dict[str, Any] | None = None,
) -> Alert:
    row = Alert(token=token, level=level.name.lower(), code=code, reason=reason, ts=ts, params=params or {})
    session.add(row)
    await session.flush()
    return row


async def open_alert_codes(session: AsyncSession, token: str) -> set[str]:
    """Codes of unacknowledged alerts — lets rules avoid raising the same alert twice."""
    rows = await session.scalars(select(Alert.code).where(Alert.token == token, Alert.acked_at.is_(None)))
    return set(rows.all())


async def list_alerts(session: AsyncSession, *, token: str | None = None, open_only: bool = False) -> Sequence[Alert]:
    stmt = select(Alert)
    if token is not None:
        stmt = stmt.where(Alert.token == token)
    if open_only:
        stmt = stmt.where(Alert.acked_at.is_(None))
    return (await session.scalars(stmt.order_by(Alert.ts.desc(), Alert.id.desc()))).all()


async def ack_alert(session: AsyncSession, *, alert_id: int, doctor_id: str, ts: datetime) -> Alert | None:
    """Doctor acknowledges an alert; the access is written to audit_log in the same transaction.

    Re-acking is a no-op (first acknowledgement wins, no duplicate audit record). Unknown id -> None.
    """
    alert = await session.get(Alert, alert_id)
    if alert is None or alert.acked_at is not None:
        return alert
    alert.acked_by = doctor_id
    alert.acked_at = ts
    await audit(
        session,
        actor=doctor_id,
        action="ack_alert",
        token=alert.token,
        details={"alert_id": alert.id, "code": alert.code, "level": alert.level},
        ts=ts,
    )
    return alert


# --- dialog -----------------------------------------------------------------------------------------


async def get_dialog_state(session: AsyncSession, token: str) -> DialogState | None:
    return await session.get(DialogState, token)


async def set_dialog_state(
    session: AsyncSession, *, token: str, state: str, data: dict[str, Any], now: datetime
) -> DialogState:
    row = await session.get(DialogState, token)
    if row is None:
        row = DialogState(token=token, state=state, data=data, updated_at=now)
        session.add(row)
    else:
        row.state = state
        row.data = dict(data)  # new object so SQLAlchemy sees the JSON change
        row.updated_at = now
    await session.flush()
    return row


async def add_message(
    session: AsyncSession,
    *,
    token: str,
    direction: str,
    text: ScrubbedText,
    ts: datetime,
    fsm_state: str | None = None,
) -> Message:
    if direction not in ("in", "out"):
        raise ValueError("direction must be 'in' or 'out'")
    row = Message(token=token, direction=direction, text_scrubbed=text, ts=ts, fsm_state=fsm_state)
    session.add(row)
    await session.flush()
    return row


# --- demo clock persistence (ADR-005) ----------------------------------------------------------------


async def load_clock_offset(session: AsyncSession) -> float:
    row = await session.get(ClockState, 1)
    return row.offset_seconds if row else 0.0


async def save_clock_offset(session: AsyncSession, offset_seconds: float) -> None:
    row = await session.get(ClockState, 1)
    if row is None:
        session.add(ClockState(id=1, offset_seconds=offset_seconds))
    else:
        row.offset_seconds = offset_seconds
    await session.flush()