vignette/apps/api/app/session_persistence.py

"""DB-backed counseling session persistence with in-process fallback support."""

from __future__ import annotations

import time
import uuid
from datetime import datetime, timezone
from typing import Any, Iterable

from .db import acquire, get_pool
from .deps import Principal
from .config import settings
from .persona_repository import SEED_VERSION, card_from_row, seed_fallback_persona, seed_persona_id
from .runtime_policy import require_runtime_fallback_allowed, runtime_fallback_allowed
from .services import memory, state_machine
from .services.persona import PersonaCard
from .store import DEFAULT_TURN_VISIBLE_TO, InProcSession, TurnRecord

_EVALUATION_CACHE: dict[str, dict[str, Any]] = {}
_SESSION_AUDIT_ROLES = {"teacher", "admin"}


_JOINED_CARD_COLUMNS = (
    "card_persona_id",
    "card_code",
    "card_version",
    "card_status",
    "card_display_name",
    "card_difficulty",
    "card_theory_target",
    "card_demographics",
    "card_presenting",
    "card_history",
    "card_big5",
    "card_resistance",
    "card_speech_style",
    "card_affect_baseline",
    "card_ccd",
    "card_dsm5_dimensional",
    "card_source_provenance",
    "card_is_synthetic",
)


def _ts(value: datetime | None) -> float | None:
    if value is None:
        return None
    if value.tzinfo is None:
        value = value.replace(tzinfo=timezone.utc)
    return value.timestamp()


def _row_value(row, key: str):
    try:
        return row[key]
    except Exception:
        return None


def _card_from_joined_session_row(row) -> PersonaCard | None:
    if _row_value(row, "card_persona_id") is None:
        return None
    card_row = {
        key.removeprefix("card_"): _row_value(row, key)
        for key in _JOINED_CARD_COLUMNS
    }
    return card_from_row(card_row)


def _stage(stage: object) -> str:
    return getattr(stage, "value", str(stage))


async def _record_session_read_audit(
    conn: Any,
    principal: Principal,
    *,
    target_kind: str,
    target_id: str,
    detail: dict[str, Any],
) -> None:
    if principal.role.value not in _SESSION_AUDIT_ROLES:
        return
    await conn.execute(
        """
        INSERT INTO audit.audit_log (
          actor_uid, action, target_kind, target_id, detail
        )
        VALUES ($1::uuid, $2, $3, $4, $5::jsonb)
        """,
        principal.user_id,
        "read_session",
        target_kind,
        target_id,
        detail,
    )


def _state_from_row(row, card: PersonaCard) -> state_machine.SessionState:
    if row is None:
        return state_machine.init_state(
            params=card.openness_params(),
        )
    return state_machine.SessionState(
        stage=state_machine.Stage(row["stage"]),
        turn_seq=int(row["turn_seq"]),
        effective_openness=float(row["effective_openness"]),
        rapport_credit=float(row["rapport_credit"]),
        resistance=float(row["resistance"]),
        ideation_stage=int(row["ideation_stage"]),
        turns_in_stage=int(row["turns_in_stage"] or 0),
        affect_state=dict(row["affect_state"] or {}),
    )


def _turn_from_row(row) -> TurnRecord:
    created_at = _ts(row["created_at"]) or time.time()
    return TurnRecord(
        turn_seq=int(row["seq"]),
        speaker=row["speaker"],
        stage=row["stage"],
        text=row["text"] or row["text_masked"] or "",
        text_masked=row["text_masked"] or row["text"] or "",
        created_at=created_at,
        llm_provider=_row_value(row, "llm_provider"),
        model=_row_value(row, "model"),
        tokens_in=_row_value(row, "tokens_in"),
        tokens_out=_row_value(row, "tokens_out"),
        cost_usd=_row_value(row, "cost_usd"),
        audio_ref=_row_value(row, "audio_ref"),
        silence_ms=_row_value(row, "silence_ms"),
        speech_rate=_row_value(row, "speech_rate"),
        barge_in=_row_value(row, "barge_in"),
        visible_to=tuple(_row_value(row, "visible_to") or DEFAULT_TURN_VISIBLE_TO),
    )


async def ensure_review_tables() -> None:
    """Create runtime review/evaluation storage when the DB role allows it."""
    try:
        get_pool()
        async with acquire(role="admin") as conn:
            await conn.execute(
                """
                CREATE TABLE IF NOT EXISTS app.session_evaluation (
                  session_id UUID PRIMARY KEY REFERENCES app.sessions(id) ON DELETE CASCADE,
                  status TEXT NOT NULL CHECK (status IN ('ready','degraded','error')),
                  source TEXT NOT NULL,
                  scope TEXT NOT NULL,
                  stage TEXT NOT NULL,
                  payload JSONB NOT NULL DEFAULT '{}'::jsonb,
                  error TEXT,
                  created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
                  updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
                )
                """
            )
            await conn.execute(
                """
                ALTER TABLE app.session_evaluation ENABLE ROW LEVEL SECURITY;

                DROP POLICY IF EXISTS p_session_evaluation_select ON app.session_evaluation;
                DROP POLICY IF EXISTS p_session_evaluation_insert ON app.session_evaluation;
                DROP POLICY IF EXISTS p_session_evaluation_update ON app.session_evaluation;
                DROP POLICY IF EXISTS p_session_evaluation_delete ON app.session_evaluation;

                CREATE POLICY p_session_evaluation_select
                  ON app.session_evaluation FOR SELECT USING (
                    app.is_ai_context()
                    OR app.current_role_name() IN ('admin','instructor')
                    OR EXISTS (
                      SELECT 1 FROM app.sessions s
                      WHERE s.id = app.session_evaluation.session_id
                        AND s.learner_id = app.current_uid()
                    )
                  );
                CREATE POLICY p_session_evaluation_insert
                  ON app.session_evaluation FOR INSERT WITH CHECK (
                    app.is_ai_context()
                    OR app.current_role_name() IN ('admin','instructor')
                    OR EXISTS (
                      SELECT 1 FROM app.sessions s
                      WHERE s.id = app.session_evaluation.session_id
                        AND s.learner_id = app.current_uid()
                    )
                  );
                CREATE POLICY p_session_evaluation_update
                  ON app.session_evaluation FOR UPDATE USING (
                    app.is_ai_context()
                    OR app.current_role_name() IN ('admin','instructor')
                    OR EXISTS (
                      SELECT 1 FROM app.sessions s
                      WHERE s.id = app.session_evaluation.session_id
                        AND s.learner_id = app.current_uid()
                    )
                  ) WITH CHECK (
                    app.is_ai_context()
                    OR app.current_role_name() IN ('admin','instructor')
                    OR EXISTS (
                      SELECT 1 FROM app.sessions s
                      WHERE s.id = app.session_evaluation.session_id
                        AND s.learner_id = app.current_uid()
                    )
                  );
                CREATE POLICY p_session_evaluation_delete
                  ON app.session_evaluation FOR DELETE USING (
                    app.is_ai_context()
                    OR app.current_role_name() IN ('admin','instructor')
                    OR EXISTS (
                      SELECT 1 FROM app.sessions s
                      WHERE s.id = app.session_evaluation.session_id
                        AND s.learner_id = app.current_uid()
                    )
                  )
                """
            )
    except Exception:
        return


async def save_session_evaluation(
    *,
    session_id: str,
    learner_id: str,
    status: str,
    source: str,
    scope: str,
    stage: str,
    payload: dict[str, Any],
    error: str | None = None,
) -> bool:
    record = {
        "status": status,
        "source": source,
        "scope": scope,
        "stage": stage,
        "payload": payload,
        "error": error,
    }
    if runtime_fallback_allowed():
        _EVALUATION_CACHE[session_id] = record
    try:
        get_pool()
        async with acquire(role="learner", user_id=learner_id) as conn:
            await conn.execute(
                """
                INSERT INTO app.session_evaluation (
                  session_id, status, source, scope, stage, payload, error,
                  created_at, updated_at
                )
                VALUES ($1::uuid, $2, $3, $4, $5, $6::jsonb, $7, now(), now())
                ON CONFLICT (session_id) DO UPDATE SET
                  status = EXCLUDED.status,
                  source = EXCLUDED.source,
                  scope = EXCLUDED.scope,
                  stage = EXCLUDED.stage,
                  payload = EXCLUDED.payload,
                  error = EXCLUDED.error,
                  updated_at = now()
                """,
                session_id,
                status,
                source,
                scope,
                stage,
                payload,
                error,
            )
        return True
    except Exception:
        require_runtime_fallback_allowed("session evaluation")
        return False


async def load_session_evaluation(
    session_id: str,
    principal: Principal,
) -> tuple[dict[str, Any] | None, bool]:
    try:
        get_pool()
        async with acquire(
            role=principal.role.value,
            user_id=principal.user_id,
            cohort_ids=principal.cohort_ids,
        ) as conn:
            row = await conn.fetchrow(
                """
                SELECT status, source, scope, stage, payload, error, updated_at
                FROM app.session_evaluation
                WHERE session_id = $1::uuid
                """,
                session_id,
            )
            if row is None:
                return (
                    _EVALUATION_CACHE.get(session_id) if runtime_fallback_allowed() else None
                ), False
            return {
                "status": row["status"],
                "source": row["source"],
                "scope": row["scope"],
                "stage": row["stage"],
                "payload": dict(row["payload"] or {}),
                "error": row["error"],
                "updated_at": _ts(row["updated_at"]),
            }, True
    except Exception:
        require_runtime_fallback_allowed("session evaluation")
        return _EVALUATION_CACHE.get(session_id), False


def _session_from_rows(row, state_row, turn_rows: Iterable) -> InProcSession | None:
    card = _card_from_joined_session_row(row)
    if card is None:
        if not settings.allow_seed_persona_fallback:
            return None
        persona_code = (row["persona_code"] or "").upper()
        legacy_entry = seed_fallback_persona(persona_code)
        if legacy_entry is None:
            return None
        card = legacy_entry.card
    persona_code = card.code

    ended_at = _ts(row["ended_at"])
    started_at = _ts(row["started_at"]) or time.time()
    return InProcSession(
        session_id=str(row["id"]),
        case_id=str(row["runtime_case_id"] or row["case_id"] or row["id"]),
        learner_id=str(row["learner_id"]),
        persona_code=persona_code,
        theory_mode=row["theory_mode"] or "humanistic",
        persona=card,
        state=_state_from_row(state_row, card),
        session_no=int(row["session_no"] or 1),
        created_at=started_at,
        ended_at=ended_at,
        turns=[_turn_from_row(turn_row) for turn_row in turn_rows],
        ended=ended_at is not None,
        prev_rapport_credit=float(row["prev_rapport_credit"] or 0.0),
    )


async def _upsert_state(conn, session_id: str, state: state_machine.SessionState) -> None:
    await conn.execute(
        """
        INSERT INTO app.session_state (
          session_id, stage, turn_seq, effective_openness, rapport_credit,
          resistance, ideation_stage, turns_in_stage, affect_state, updated_at
        )
        VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9::jsonb, now())
        ON CONFLICT (session_id) DO UPDATE SET
          stage = EXCLUDED.stage,
          turn_seq = EXCLUDED.turn_seq,
          effective_openness = EXCLUDED.effective_openness,
          rapport_credit = EXCLUDED.rapport_credit,
          resistance = EXCLUDED.resistance,
          ideation_stage = EXCLUDED.ideation_stage,
          turns_in_stage = EXCLUDED.turns_in_stage,
          affect_state = EXCLUDED.affect_state,
          updated_at = now()
        """,
        session_id,
        _stage(state.stage),
        state.turn_seq,
        state.effective_openness,
        state.rapport_credit,
        state.resistance,
        state.ideation_stage,
        state.turns_in_stage,
        state.affect_state,
    )


async def create_session(
    *,
    learner_id: str,
    card: PersonaCard,
    theory_mode: str,
    state: state_machine.SessionState,
    session_no: int = 1,
    carry_rapport: float = 0.0,
    persona_id: str | None = None,
    persona_version: int | None = None,
) -> InProcSession | None:
    """Create a DB-backed session, returning None when DB persistence is unavailable."""
    try:
        get_pool()
        runtime_case_id = str(uuid.uuid4())
        pinned_persona_id = persona_id or seed_persona_id(card.code)
        pinned_persona_version = persona_version or SEED_VERSION
        async with acquire(role="learner", user_id=learner_id) as conn:
            row = await conn.fetchrow(
                """
                INSERT INTO app.sessions (
                  runtime_case_id, learner_id, persona_id, persona_version,
                  persona_code, persona_display_name, persona_difficulty,
                  session_no, theory_mode, stage_path, prev_rapport_credit
                )
                VALUES (
                  $1::uuid, $2::uuid, $3::uuid, $4,
                  $5, $6, $7,
                  $8, $9, '[]'::jsonb, $10
                )
                RETURNING id, runtime_case_id, case_id, learner_id, persona_code,
                          session_no, theory_mode, started_at, ended_at, prev_rapport_credit
                """,
                runtime_case_id,
                learner_id,
                pinned_persona_id,
                pinned_persona_version,
                card.code,
                card.display_name,
                card.difficulty,
                session_no,
                theory_mode,
                carry_rapport,
            )
            await _upsert_state(conn, str(row["id"]), state)
        return InProcSession(
            session_id=str(row["id"]),
            case_id=runtime_case_id,
            learner_id=learner_id,
            persona_code=card.code,
            theory_mode=theory_mode,
            persona=card,
            state=state,
            session_no=session_no,
            created_at=_ts(row["started_at"]) or time.time(),
            ended_at=None,
            turns=[],
            ended=False,
            prev_rapport_credit=carry_rapport,
        )
    except Exception:
        require_runtime_fallback_allowed("session creation")
        return None


async def load_session(
    session_id: str,
    principal: Principal,
    *,
    allow_ended: bool = False,
) -> InProcSession | None:
    try:
        get_pool()
        async with acquire(
            role=principal.role.value,
            user_id=principal.user_id,
            cohort_ids=principal.cohort_ids,
        ) as conn:
            row = await conn.fetchrow(
                """
                SELECT
                  s.id, s.runtime_case_id, s.case_id, s.learner_id, s.persona_code,
                  s.session_no, s.theory_mode, s.started_at, s.ended_at,
                  s.prev_rapport_credit,
                  pc.persona_id AS card_persona_id,
                  pc.code AS card_code,
                  pc.version AS card_version,
                  pc.status AS card_status,
                  pc.display_name AS card_display_name,
                  pc.difficulty AS card_difficulty,
                  pc.theory_target AS card_theory_target,
                  pc.demographics AS card_demographics,
                  pc.presenting AS card_presenting,
                  pc.history AS card_history,
                  pc.big5 AS card_big5,
                  pc.resistance AS card_resistance,
                  pc.speech_style AS card_speech_style,
                  pc.affect_baseline AS card_affect_baseline,
                  pc.ccd AS card_ccd,
                  pc.dsm5_dimensional AS card_dsm5_dimensional,
                  pc.source_provenance AS card_source_provenance,
                  pc.is_synthetic AS card_is_synthetic
                FROM app.sessions s
                LEFT JOIN app.persona_card pc
                  ON pc.persona_id = s.persona_id
                 AND pc.version = s.persona_version
                WHERE s.id = $1::uuid
                """,
                session_id,
            )
            if row is None:
                return None
            if row["ended_at"] is not None and not allow_ended:
                return None
            state_row = await conn.fetchrow(
                """
                SELECT stage, turn_seq, effective_openness, rapport_credit, resistance,
                       ideation_stage, turns_in_stage, affect_state
                FROM app.session_state
                WHERE session_id = $1::uuid
                """,
                session_id,
            )
            turn_rows = await conn.fetch(
                """
                SELECT seq, speaker, stage, text, text_masked, created_at,
                       llm_provider, model, tokens_in, tokens_out, cost_usd,
                       audio_ref, silence_ms, speech_rate, barge_in, visible_to
                FROM app.turns
                WHERE session_id = $1::uuid
                ORDER BY seq
                """,
                session_id,
            )
            sess = _session_from_rows(row, state_row, turn_rows)
            if sess is not None:
                await _record_session_read_audit(
                    conn,
                    principal,
                    target_kind="session",
                    target_id=session_id,
                    detail={
                        "access": "load_session",
                        "role": principal.role.value,
                        "learner_id": sess.learner_id,
                    },
                )
            return sess
    except Exception:
        require_runtime_fallback_allowed("session load")
        return None


async def append_turn(
    *,
    session_id: str,
    learner_id: str,
    turn: TurnRecord,
) -> bool:
    try:
        get_pool()
        async with acquire(role="learner", user_id=learner_id) as conn:
            locked = await conn.fetchval(
                "SELECT id FROM app.sessions WHERE id = $1::uuid FOR UPDATE",
                session_id,
            )
            if locked is None:
                return False
            seq = int(
                await conn.fetchval(
                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM app.turns WHERE session_id = $1::uuid",
                    session_id,
                )
                or 1
            )
            await conn.execute(
                """
                INSERT INTO app.turns (
                  session_id, seq, speaker, stage, text, text_masked, actor_kind,
                  llm_provider, model, tokens_in, tokens_out, cost_usd,
                  audio_ref, silence_ms, speech_rate, barge_in, visible_to
                )
                VALUES (
                  $1::uuid, $2, $3, $4, $5, $6, $7,
                  $8, $9, $10, $11, $12,
                  $13, $14, $15, $16, $17::text[]
                )
                ON CONFLICT (session_id, seq) DO NOTHING
                """,
                session_id,
                seq,
                turn.speaker,
                turn.stage,
                turn.text_masked,
                turn.text_masked,
                "human_learner" if turn.speaker == "counselor" else "client_ai",
                turn.llm_provider,
                turn.model,
                turn.tokens_in,
                turn.tokens_out,
                turn.cost_usd,
                turn.audio_ref,
                turn.silence_ms,
                turn.speech_rate,
                turn.barge_in,
                list(turn.visible_to or DEFAULT_TURN_VISIBLE_TO),
            )
        return True
    except Exception:
        require_runtime_fallback_allowed("session turn append")
        return False


async def update_state(
    *,
    session_id: str,
    learner_id: str,
    state: state_machine.SessionState,
) -> bool:
    try:
        get_pool()
        async with acquire(role="learner", user_id=learner_id) as conn:
            await _upsert_state(conn, session_id, state)
        return True
    except Exception:
        require_runtime_fallback_allowed("session state update")
        return False


async def end_session(sess: InProcSession, carry: memory.CarryOver) -> bool:
    try:
        get_pool()
        digest = (
            f"회기 축어록 {len(sess.turns)}개가 저장되었습니다. 정밀 리뷰는 생성 대기 중입니다."
            if sess.turns
            else "실제 발화가 없어 요약을 생성하지 않았습니다."
        )
        async with acquire(role="learner", user_id=sess.learner_id) as conn:
            await conn.execute(
                """
                UPDATE app.sessions
                SET ended_at = COALESCE(ended_at, now())
                WHERE id = $1::uuid
                """,
                sess.session_id,
            )
            await _upsert_state(conn, sess.session_id, sess.state)
            await conn.execute(
                """
                INSERT INTO app.session_summary (
                  session_id, case_id, session_no, end_state, rapport_delta,
                  digest, open_threads, created_at
                )
                VALUES ($1::uuid, $2::uuid, $3, $4::jsonb, $5, $6, $7::jsonb, now())
                ON CONFLICT (session_id) DO UPDATE SET
                  end_state = EXCLUDED.end_state,
                  rapport_delta = EXCLUDED.rapport_delta,
                  digest = EXCLUDED.digest,
                  open_threads = EXCLUDED.open_threads
                """,
                sess.session_id,
                sess.case_id,
                sess.session_no,
                carry.end_state,
                carry.rapport_delta,
                digest,
                list(carry.compression_job.open_threads if carry.compression_job else []),
            )
        return True
    except Exception:
        require_runtime_fallback_allowed("session end")
        return False


async def list_sessions(principal: Principal) -> tuple[list[InProcSession], bool]:
    try:
        get_pool()
        learner_filter = "WHERE s.learner_id = $1::uuid" if principal.role.value == "learner" else ""
        query_args = [principal.user_id] if principal.role.value == "learner" else []
        async with acquire(
            role=principal.role.value,
            user_id=principal.user_id,
            cohort_ids=principal.cohort_ids,
        ) as conn:
            rows = await conn.fetch(
                f"""
                SELECT
                  s.id, s.runtime_case_id, s.case_id, s.learner_id, s.persona_code,
                  s.session_no, s.theory_mode, s.started_at, s.ended_at,
                  s.prev_rapport_credit,
                  pc.persona_id AS card_persona_id,
                  pc.code AS card_code,
                  pc.version AS card_version,
                  pc.status AS card_status,
                  pc.display_name AS card_display_name,
                  pc.difficulty AS card_difficulty,
                  pc.theory_target AS card_theory_target,
                  pc.demographics AS card_demographics,
                  pc.presenting AS card_presenting,
                  pc.history AS card_history,
                  pc.big5 AS card_big5,
                  pc.resistance AS card_resistance,
                  pc.speech_style AS card_speech_style,
                  pc.affect_baseline AS card_affect_baseline,
                  pc.ccd AS card_ccd,
                  pc.dsm5_dimensional AS card_dsm5_dimensional,
                  pc.source_provenance AS card_source_provenance,
                  pc.is_synthetic AS card_is_synthetic
                FROM app.sessions s
                LEFT JOIN app.persona_card pc
                  ON pc.persona_id = s.persona_id
                 AND pc.version = s.persona_version
                {learner_filter}
                ORDER BY s.started_at DESC
                LIMIT 100
                """,
                *query_args,
            )
            sessions: list[InProcSession] = []
            for row in rows:
                session_id = str(row["id"])
                state_row = await conn.fetchrow(
                    """
                    SELECT stage, turn_seq, effective_openness, rapport_credit, resistance,
                           ideation_stage, turns_in_stage, affect_state
                    FROM app.session_state
                    WHERE session_id = $1::uuid
                    """,
                    session_id,
                )
                turn_rows = await conn.fetch(
                    """
                    SELECT seq, speaker, stage, text, text_masked, created_at,
                           llm_provider, model, tokens_in, tokens_out, cost_usd,
                           audio_ref, silence_ms, speech_rate, barge_in, visible_to
                    FROM app.turns
                    WHERE session_id = $1::uuid
                    ORDER BY seq
                    """,
                    session_id,
                )
                sess = _session_from_rows(row, state_row, turn_rows)
                if sess is not None:
                    sessions.append(sess)
            await _record_session_read_audit(
                conn,
                principal,
                target_kind="session_list",
                target_id="sessions",
                detail={
                    "access": "list_sessions",
                    "role": principal.role.value,
                    "result_count": len(sessions),
                },
            )
            return sessions, True
    except Exception:
        require_runtime_fallback_allowed("session list")
        return [], False