vignette/apps/api/app/routes/voice.py

"""Voice routes for the OpenAI STT/TTS cascade over WebSocket.

Client sends JSON controls plus binary audio chunks:
  audio_start -> binary audio chunks -> audio_end

Server emits:
  ready -> state(listening) -> state(thinking) -> transcript -> reply
  -> state(speaking) -> tts_chunk + binary audio chunks -> tts_end -> state(idle)

When voice is not configured, the route reports degraded state and closes
cleanly instead of crashing.
"""

from __future__ import annotations

import json
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from starlette.websockets import WebSocketState

from .. import session_persistence
from ..auth_sessions import get_session
from ..config import settings
from ..deps import Principal, Role
from ..engine_client import EngineError, engine_client
from ..persona_repository import get_catalog_persona
from ..runtime_policy import require_runtime_fallback_allowed, runtime_fallback_allowed
from ..services import memory, orchestrator, state_machine
from ..services import voice as voice_svc
from ..services.voice import VoicePreset, VoiceUnavailable, resolve_voice, voice_service
from ..store import InProcSession, TurnRecord, store

router = APIRouter(prefix="/voice", tags=["voice"])

# WebSocket close codes.
WS_CLOSE_DEGRADED = 1011
WS_CLOSE_BAD_REQUEST = 1008
WS_CLOSE_UNAUTHORIZED = 1008

# Per-utterance audio cap to avoid unbounded memory growth.
_MAX_AUDIO_BYTES = 10 * 1024 * 1024


@router.get("/health")
async def voice_health() -> JSONResponse:
    """Return voice service readiness."""
    available = voice_service.is_available()
    body = {
        "status": "ok" if available else "degraded",
        "available": available,
        "stt_model": voice_svc.STT_MODEL,
        "tts_model": voice_svc.TTS_MODEL,
        "reason": None if available else "OPENAI_API_KEY is not configured",
    }
    return JSONResponse(body, status_code=200 if available else 503)


@router.websocket("/ws")
async def voice_ws(websocket: WebSocket) -> None:
    """Run one authenticated learner voice cascade."""
    await websocket.accept()

    # Authenticate the same server-side browser session used by REST routes.
    principal = await _principal_from_websocket(websocket)
    if principal is None:
        await _safe_send_json(websocket, {"type": "error", "detail": "not authenticated"})
        await _safe_close(websocket, WS_CLOSE_UNAUTHORIZED)
        return
    if principal.role != Role.LEARNER:
        await _safe_send_json(websocket, {"type": "error", "detail": "only learners can use voice"})
        await _safe_close(websocket, WS_CLOSE_UNAUTHORIZED)
        return

    # Bind to an existing session first. persona_code creation is dev-only.
    session_id, voice_preset, err, bind_meta = await _bind_session(websocket, principal)
    if err is not None:
        await _safe_send_json(websocket, {"type": "error", "detail": err})
        await _safe_close(websocket, WS_CLOSE_BAD_REQUEST)
        return
    assert session_id is not None and voice_preset is not None
    if bind_meta.get("degraded"):
        await _safe_send_json(
            websocket,
            {
                "type": "degraded",
                "reason": bind_meta.get("degraded_reason", "voice session binding degraded"),
                **bind_meta,
            },
        )

    # Voice misconfiguration is reported explicitly and then closed cleanly.
    if not voice_service.is_available():
        await _safe_send_json(
            websocket,
            {"type": "degraded", "reason": "OPENAI_API_KEY is not configured"},
        )
        await _safe_close(websocket, WS_CLOSE_DEGRADED)
        return

    await _safe_send_json(
        websocket,
        {
            "type": "ready",
            "session_id": session_id,
            "voice": voice_preset.openai_voice,
            "preset": voice_preset.preset,
            "state": "idle",
            **bind_meta,
        },
    )

    audio_buf = bytearray()
    receiving = False

    try:
        while True:
            msg = await websocket.receive()
            mtype = msg.get("type")
            if mtype == "websocket.disconnect":
                break

            # Binary frames are audio chunks.
            if msg.get("bytes") is not None:
                if not receiving:
                    # Be tolerant when audio arrives before audio_start.
                    receiving = True
                    audio_buf.clear()
                    await _safe_send_json(websocket, {"type": "state", "state": "listening"})
                audio_buf.extend(msg["bytes"])
                if len(audio_buf) > _MAX_AUDIO_BYTES:
                    await _safe_send_json(
                        websocket,
                        {"type": "error", "detail": "audio too large; please send a shorter utterance"},
                    )
                    audio_buf.clear()
                    receiving = False
                continue

            # Text frames are JSON controls.
            text = msg.get("text")
            if text is None:
                continue
            try:
                ctrl = json.loads(text)
            except (json.JSONDecodeError, TypeError):
                await _safe_send_json(websocket, {"type": "error", "detail": "invalid control json"})
                continue

            ctype = ctrl.get("type")
            if ctype == "audio_start":
                receiving = True
                audio_buf.clear()
                await _safe_send_json(websocket, {"type": "state", "state": "listening"})

            elif ctype == "audio_end":
                receiving = False
                await _handle_utterance(
                    websocket,
                    session_id=session_id,
                    principal=principal,
                    voice_preset=voice_preset,
                    audio=bytes(audio_buf),
                    fmt=ctrl.get("format"),
                )
                audio_buf.clear()

            elif ctype == "text_turn":
                # Text-only path for accessibility and deterministic tests.
                receiving = False
                audio_buf.clear()
                learner_text = (ctrl.get("text") or "").strip()
                if learner_text:
                    await _run_turn_and_speak(
                        websocket,
                        session_id=session_id,
                        principal=principal,
                        voice_preset=voice_preset,
                        learner_text=learner_text,
                    )

            elif ctype == "ping":
                await _safe_send_json(websocket, {"type": "pong"})

            elif ctype == "close":
                break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        await _safe_send_json(websocket, {"type": "error", "detail": f"voice ws error: {e}"})
    finally:
        await _safe_close(websocket)


async def _handle_utterance(
    websocket: WebSocket,
    *,
    session_id: str,
    principal: Principal,
    voice_preset: VoicePreset,
    audio: bytes,
    fmt: Optional[str],
) -> None:
    """Transcribe one utterance, generate the client reply, then synthesize TTS."""
    if not audio:
        await _safe_send_json(websocket, {"type": "transcript", "text": "", "final": True})
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    # STT begins after the learner stops speaking.
    await _safe_send_json(websocket, {"type": "state", "state": "thinking"})
    filename, content_type = _audio_meta(fmt)
    try:
        stt = await voice_service.transcribe(
            audio, filename=filename, content_type=content_type
        )
    except VoiceUnavailable as e:
        await _safe_send_json(websocket, {"type": "degraded", "reason": str(e)})
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return
    except Exception as e:
        await _safe_send_json(websocket, {"type": "error", "detail": f"STT failed: {e}"})
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    learner_text = stt.text
    await _safe_send_json(
        websocket,
        {"type": "transcript", "text": learner_text, "final": True, "speaker": "counselor"},
    )
    if not learner_text:
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    await _run_turn_and_speak(
        websocket,
        session_id=session_id,
        principal=principal,
        voice_preset=voice_preset,
        learner_text=learner_text,
    )


async def _run_turn_and_speak(
    websocket: WebSocket,
    *,
    session_id: str,
    principal: Principal,
    voice_preset: VoicePreset,
    learner_text: str,
) -> None:
    """Run one counseling turn and stream synthesized client speech."""
    sess, err = await _load_voice_session(session_id, principal)
    if sess is None:
        await _safe_send_json(websocket, {"type": "error", "detail": err or "session not found or ended"})
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    recall = memory.RecallContext()
    ctx = orchestrator.prepare_turn(
        session_id=session_id,
        case_id=sess.case_id,
        card=sess.persona,
        state=sess.state,
        learner_text=learner_text,
        recall_summary=recall.recall_summary,
        pinned_facts=recall.pinned_facts,
        recent_turns=sess.recent_turns(),
    )
    assert ctx.state_after is not None

    # Voice needs the full client reply before TTS starts.
    try:
        result = await orchestrator.run_turn_generate(ctx, engine_client)
    except EngineError as e:
        await _safe_send_json(websocket, {"type": "error", "detail": f"engine unavailable: {e}"})
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    reply = result.client_reply or ""
    # Persist only after the client reply has been generated. A failed AI turn
    # must not leave a learner-only transcript in review or history.
    await _append_voice_turn(
        sess,
        TurnRecord(
            turn_seq=ctx.state_after.turn_seq,
            speaker="counselor",
            stage=ctx.state_after.stage.value,
            text=learner_text,
            text_masked=ctx.learner_text_masked,
        ),
    )
    if reply:
        # Persist the generated client reply before TTS playback.
        await _append_voice_turn(
            sess,
            TurnRecord(
                turn_seq=result.turn_seq,
                speaker="client",
                stage=result.stage,
                text=reply,
                text_masked=reply,
            ),
        )
    await _update_voice_state(sess, result.state_after)

    # Send the final client text before audio playback.
    await _safe_send_json(
        websocket,
        {
            "type": "reply",
            "text": reply,
            "speaker": "client",
            "stage": result.stage,
            "effective_openness": round(result.effective_openness, 4),
            "turn_seq": result.turn_seq,
            "safety_flagged": result.safety_flagged,
            "crisis_kind": result.crisis_kind,
        },
    )

    if not reply:
        await _safe_send_json(websocket, {"type": "state", "state": "idle"})
        return

    # TTS speaking state comes before chunk metadata and binary audio.
    await _safe_send_json(
        websocket,
        {"type": "state", "state": "speaking", "voice": voice_preset.openai_voice},
    )
    try:
        n = 0
        async for ck in voice_service.synthesize_stream(reply, voice_preset):
            # Metadata precedes the binary chunk so the client can pair them.
            await _safe_send_json(
                websocket, {"type": "tts_chunk", "seq": ck.seq, "rms": round(ck.rms, 4)}
            )
            await _safe_send_bytes(websocket, ck.audio)
            n += 1
        await _safe_send_json(websocket, {"type": "tts_end", "chunks": n})
    except VoiceUnavailable as e:
        await _safe_send_json(websocket, {"type": "degraded", "reason": str(e)})
    except Exception as e:
        await _safe_send_json(websocket, {"type": "error", "detail": f"TTS failed: {e}"})

    await _safe_send_json(websocket, {"type": "state", "state": "idle"})


async def _load_voice_session(
    session_id: str,
    principal: Principal,
) -> tuple[InProcSession | None, str | None]:
    sess = await session_persistence.load_session(session_id, principal, allow_ended=True)
    if sess is not None:
        store.put(sess)
    elif runtime_fallback_allowed():
        sess = store.get(session_id)
    if sess is None:
        return None, f"unknown session {session_id}"
    if sess.learner_id != principal.user_id:
        return None, "session does not belong to user"
    if sess.ended:
        return None, "session already ended"
    return sess, None


async def _append_voice_turn(sess: InProcSession, turn: TurnRecord) -> None:
    if await session_persistence.append_turn(
        session_id=sess.session_id,
        learner_id=sess.learner_id,
        turn=turn,
    ):
        sess.turns.append(turn)
        store.put(sess)
        return
    require_runtime_fallback_allowed("voice session turn append")
    store.append_turn(sess.session_id, turn)


async def _update_voice_state(
    sess: InProcSession,
    state: state_machine.SessionState,
) -> None:
    if await session_persistence.update_state(
        session_id=sess.session_id,
        learner_id=sess.learner_id,
        state=state,
    ):
        sess.state = state
        store.put(sess)
        return
    require_runtime_fallback_allowed("voice session state update")
    store.update_state(sess.session_id, state)


async def _principal_from_websocket(websocket: WebSocket) -> Principal | None:
    """Restore the same server-side browser session used by REST routes."""
    raw_cookie = websocket.cookies.get(settings.cookie_name)
    if raw_cookie is None and settings.environment == "dev":
        raw_cookie = websocket.cookies.get("vignette_sid")

    session = await get_session(raw_cookie)
    if session is None:
        return None

    try:
        role = Role(session.role)
    except ValueError:
        return None

    return Principal(
        user_id=session.user_id,
        role=role,
        cohort_ids=session.cohort_ids,
        email=session.email,
        display_name=session.display_name,
    )


async def _bind_session(
    websocket: WebSocket,
    principal: Principal,
) -> tuple[Optional[str], Optional[VoicePreset], Optional[str], dict[str, object]]:
    """Bind an existing session or create a dev-only voice session."""
    qp = websocket.query_params
    explicit_preset = qp.get("preset")

    session_id = qp.get("session_id")
    if session_id:
        sess, err = await _load_voice_session(session_id, principal)
        if sess is None:
            return None, None, err or f"unknown session {session_id}", {}
        vp = resolve_voice(persona_code=sess.persona.code, preset=explicit_preset)
        return session_id, vp, None, {"degraded": False, "persona_catalog_source": "session"}

    # persona_code session creation is local-dev only. Production uses REST start.
    if settings.environment != "dev":
        return None, None, "session_id required", {}

    persona_code = qp.get("persona_code")
    if not persona_code:
        return None, None, "session_id or persona_code query required", {}
    try:
        catalog_persona = await get_catalog_persona(persona_code)
    except Exception:
        return None, None, "persona catalog database unavailable", {}
    if catalog_persona is None:
        return None, None, f"unknown persona {persona_code}", {}
    card = catalog_persona.card

    st = state_machine.init_state(
        base_resistance=card.base_resistance(),
        unlock_rate=card.unlock_rate(),
        decay_floor=card.decay_floor(),
        ideation_baseline=card.ideation_baseline(),
    )
    sess = await session_persistence.create_session(
        learner_id=principal.user_id,
        card=card,
        theory_mode="humanistic",
        state=st,
        session_no=1,
        carry_rapport=st.rapport_credit,
        persona_id=catalog_persona.persona_id,
        persona_version=catalog_persona.version,
    )
    session_source = "database"
    if sess is None:
        require_runtime_fallback_allowed("voice session creation")
        sess = store.create(
            learner_id=principal.user_id,
            persona=card,
            theory_mode="humanistic",
            state=st,
            session_no=1,
            carry_rapport=st.rapport_credit,
        )
        session_source = "runtime"
    else:
        store.put(sess)
    vp = resolve_voice(persona_code=card.code, preset=explicit_preset)
    degraded_reasons: list[str] = []
    if catalog_persona.degraded:
        degraded_reasons.append("카탈로그 원본을 확인하지 못해 음성 회기를 시작하지 않습니다")
    if session_source == "runtime":
        degraded_reasons.append("세션 저장소 연결 전까지 비영구 개발 런타임 기록을 사용합니다")
    bind_meta = {
        "degraded": bool(degraded_reasons),
        "degraded_reason": "; ".join(degraded_reasons) if degraded_reasons else None,
        "persona_catalog_source": catalog_persona.source,
        "session_source": session_source,
    }
    return sess.session_id, vp, None, bind_meta


def _audio_meta(fmt: Optional[str]) -> tuple[str, str]:
    """Map the browser audio format to upload metadata."""
    f = (fmt or "webm").lower().lstrip(".")
    table = {
        "webm": ("audio.webm", "audio/webm"),
        "ogg": ("audio.ogg", "audio/ogg"),
        "opus": ("audio.ogg", "audio/ogg"),
        "wav": ("audio.wav", "audio/wav"),
        "mp3": ("audio.mp3", "audio/mpeg"),
        "mp4": ("audio.mp4", "audio/mp4"),
        "m4a": ("audio.m4a", "audio/mp4"),
        "pcm": ("audio.wav", "audio/wav"),
    }
    return table.get(f, ("audio.webm", "audio/webm"))


async def _safe_send_json(websocket: WebSocket, payload: dict) -> None:
    if websocket.client_state != WebSocketState.CONNECTED:
        return
    try:
        await websocket.send_text(json.dumps(payload, ensure_ascii=False))
    except Exception:
        pass


async def _safe_send_bytes(websocket: WebSocket, data: bytes) -> None:
    if websocket.client_state != WebSocketState.CONNECTED:
        return
    try:
        await websocket.send_bytes(data)
    except Exception:
        pass


async def _safe_close(websocket: WebSocket, code: int = 1000) -> None:
    if websocket.client_state == WebSocketState.DISCONNECTED:
        return
    try:
        await websocket.close(code=code)
    except Exception:
        pass


__all__ = ["router"]