src/time_machine/ui/realtime.py

from __future__ import annotations

import asyncio
import base64
import io
import json
import wave
from pathlib import Path
from typing import Any

from fastapi import WebSocket, WebSocketDisconnect
from pydantic import BaseModel

from time_machine.application.container import AppContainer
from time_machine.domain.events import (
    AudioEvent,
    ConversationTextEvent,
    ErrorEvent,
    TemporalEvent,
    TranscriptEvent,
)


DEFAULT_SAMPLE_RATE = 16000
DEFAULT_CHANNELS = 1
MAX_TURN_AUDIO_BYTES = 16 * 1024 * 1024


async def realtime_voice_socket(websocket: WebSocket, container: AppContainer) -> None:
    await websocket.accept()
    state = _RealtimeConnectionState()
    await _send_json(websocket, {"type": "ready", "sample_rate": DEFAULT_SAMPLE_RATE})

    try:
        while True:
            message = await websocket.receive()
            if message.get("type") == "websocket.disconnect":
                return
            if "text" in message and message["text"] is not None:
                await _handle_control_message(websocket, container, state, message["text"])
            elif "bytes" in message and message["bytes"] is not None:
                state.append_audio(message["bytes"])
    except WebSocketDisconnect:
        return


async def _handle_control_message(
    websocket: WebSocket,
    container: AppContainer,
    state: "_RealtimeConnectionState",
    raw_message: str,
) -> None:
    try:
        message = json.loads(raw_message)
    except json.JSONDecodeError:
        await _send_json(websocket, {"type": "error", "message": "Invalid realtime message."})
        return

    message_type = message.get("type")
    if message_type == "start":
        encounter_id = _clean_string(message.get("encounter_id"))
        if not encounter_id:
            await _send_json(websocket, {"type": "error", "message": "Launch the machine first."})
            return
        state.encounter_id = encounter_id
        state.sample_rate = _positive_int(message.get("sample_rate"), DEFAULT_SAMPLE_RATE)
        state.channels = _positive_int(message.get("channels"), DEFAULT_CHANNELS)
        state.reset_audio()
        await _send_json(websocket, {"type": "listening", "encounter_id": encounter_id})
        return

    if message_type == "cancel":
        state.reset_audio()
        await _send_json(websocket, {"type": "listening", "encounter_id": state.encounter_id})
        return

    if message_type == "commit":
        if not state.encounter_id:
            await _send_json(websocket, {"type": "error", "message": "Launch the machine first."})
            state.reset_audio()
            return
        audio = state.consume_wav()
        if not audio:
            await _send_json(websocket, {"type": "listening", "encounter_id": state.encounter_id})
            return
        await _send_json(websocket, {"type": "processing", "message": "Decoding signal..."})
        try:
            await _stream_audio_turn(websocket, container, state.encounter_id, audio)
        except Exception as exc:
            await _send_json(websocket, {"type": "error", "message": str(exc)})
        await _send_json(websocket, {"type": "listening", "encounter_id": state.encounter_id})
        return

    await _send_json(websocket, {"type": "error", "message": f"Unknown realtime message: {message_type}"})


async def _stream_audio_turn(
    websocket: WebSocket,
    container: AppContainer,
    encounter_id: str,
    audio: bytes,
) -> None:
    iterator = container.speech_orchestrator.handle_audio_turn(
        encounter_id=encounter_id,
        audio=audio,
    )
    sentinel = object()
    while True:
        event = await asyncio.to_thread(next, iterator, sentinel)
        if event is sentinel:
            return
        await _send_event(websocket, event)


async def _send_event(websocket: WebSocket, event: TemporalEvent) -> None:
    message = _event_payload(event)
    await _send_json(websocket, message)


def _event_payload(event: TemporalEvent) -> dict[str, Any]:
    payload = {
        "type": "event",
        "event_type": event.event_type,
        "message": event.message,
        "metadata": event.metadata,
    }
    if isinstance(event, TranscriptEvent):
        payload.update(
            {
                "type": "transcript",
                "text": event.transcript.text,
                "is_final": event.transcript.is_final,
                "confidence": event.transcript.confidence,
                "language": event.transcript.language,
            }
        )
    elif isinstance(event, ConversationTextEvent):
        payload.update(
            {
                "type": "conversation_text",
                "text": event.text,
                "is_final": event.is_final,
            }
        )
    elif isinstance(event, AudioEvent):
        payload.update(_audio_payload(event))
    elif isinstance(event, ErrorEvent):
        payload.update({"type": "error", "severity": event.severity})
    else:
        serialized = event.model_dump(mode="json")
        serialized.pop("created_at", None)
        payload["event"] = serialized
    return payload


def _audio_payload(event: AudioEvent) -> dict[str, Any]:
    audio = event.audio
    path = Path(audio.path) if audio.path else None
    if path and path.exists():
        audio_b64 = base64.b64encode(path.read_bytes()).decode("ascii")
    else:
        audio_b64 = ""
    return {
        "type": "audio",
        "audio_b64": audio_b64,
        "mime_type": audio.mime_type,
        "duration_seconds": audio.duration_seconds,
        "description": audio.description,
    }


async def _send_json(websocket: WebSocket, payload: dict[str, Any]) -> None:
    await websocket.send_text(json.dumps(payload, ensure_ascii=True))


class _RealtimeConnectionState(BaseModel):
    encounter_id: str | None = None
    sample_rate: int = DEFAULT_SAMPLE_RATE
    channels: int = DEFAULT_CHANNELS
    audio_chunks: list[bytes] = []
    audio_bytes: int = 0

    def append_audio(self, chunk: bytes) -> None:
        if not chunk:
            return
        if self.audio_bytes + len(chunk) > MAX_TURN_AUDIO_BYTES:
            self.reset_audio()
            return
        self.audio_chunks.append(chunk)
        self.audio_bytes += len(chunk)

    def consume_wav(self) -> bytes:
        if not self.audio_chunks:
            return b""
        pcm = b"".join(self.audio_chunks)
        self.reset_audio()
        if not pcm:
            return b""
        return _pcm16_to_wav(pcm, self.sample_rate, self.channels)

    def reset_audio(self) -> None:
        self.audio_chunks = []
        self.audio_bytes = 0


def _pcm16_to_wav(pcm: bytes, sample_rate: int, channels: int) -> bytes:
    output = io.BytesIO()
    with wave.open(output, "wb") as handle:
        handle.setnchannels(channels)
        handle.setsampwidth(2)
        handle.setframerate(sample_rate)
        handle.writeframes(pcm)
    return output.getvalue()


def _clean_string(value: object) -> str | None:
    return value.strip() if isinstance(value, str) and value.strip() else None


def _positive_int(value: object, default: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        return default
    return parsed if parsed > 0 else default