"""WebSocket handler for streaming (pseudo-real-time) voice transcription.

Protocol (bidirectional, one session = one utterance):

Client → server:
  - Binary frames: raw audio chunks from MediaRecorder, in the order
    produced. The first chunk MUST contain the container header (WebM
    EBML / MP4 moov atom) so that the accumulated buffer decodes as a
    valid media file — this is how MediaRecorder.start(timeslice) emits
    data and we rely on it.
  - Text frame ``{"type": "done"}`` — client stopped recording. Server
    runs one final transcription on the full buffer and sends a
    ``final`` message, then closes.
  - Text frame ``{"type": "cancel"}`` — discard buffer without
    transcribing.

Server → client:
  - ``{"type": "ready"}`` — after accept, indicates the handler is
    ready to receive audio.
  - ``{"type": "interim", "text": "..."}`` — partial transcription
    computed from the buffer so far. Sent at most every ~1.2s.
  - ``{"type": "final",   "text": "..."}`` — last transcription after
    ``done`` was received. Always the last server message before close.
  - ``{"type": "error",   "detail": "..."}`` — fatal error; the socket
    is closed after sending.

Design:

- faster-whisper has NO streaming API — it's an utterance model. We
  simulate streaming by re-running the model on the GROWING buffer
  every ~1.2s. This wastes CPU (the head of the buffer is transcribed
  repeatedly) but is simple and produces monotonically improving
  interim text. For utterances under 30s (the realistic chat-input
  case) the waste is acceptable.
- We serialise transcription calls with an asyncio.Lock so multiple
  concurrent scheduler ticks don't stack up on top of each other.
- If the model isn't installed we send one error and close — there's
  no cloud fallback in the streaming path (OpenAI's API isn't cheap
  or fast enough to call once per second).
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import time

from fastapi import WebSocket, WebSocketDisconnect


class suppress_ws_errors(contextlib.suppress):
    """Suppress common errors when talking to an already-closed WebSocket."""

    def __init__(self) -> None:
        super().__init__(Exception)

log = logging.getLogger(__name__)

# Minimum wall-clock gap between interim model calls. Lower = fresher
# interim text but more CPU burn. 0.6s feels genuinely live during normal
# dictation; the asyncio.Lock guarantees we never stack ticks even if a
# single transcription pass takes longer than the interval.
_INTERIM_INTERVAL = 0.6

# Hard cap on buffer size so a runaway client can't exhaust RAM.
_MAX_BUFFER_BYTES = 40 * 1024 * 1024  # 40 MB ~ ~60 min of Opus


async def ws_transcribe_handler(websocket: WebSocket) -> None:
    """Handle a streaming-transcription WebSocket session."""
    await websocket.accept()

    # Lazy import to avoid pulling faster-whisper at module load time
    try:
        from agntspace.llm.transcribe import transcribe_buffer_sync
    except Exception as exc:  # pragma: no cover
        log.exception("Failed to import transcribe module")
        await websocket.send_json({"type": "error", "detail": f"Import failed: {exc!s}"})
        await websocket.close()
        return

    # Sanity-check that a local backend is installed. Cheap — faster_whisper's
    # top-level module has no side effects on import.
    try:
        import faster_whisper  # noqa: F401
    except ImportError:
        await websocket.send_json({
            "type": "error",
            "detail": (
                "Streaming transcription requires a local Whisper backend. "
                "Install it with `pip install faster-whisper`."
            ),
        })
        await websocket.close()
        return

    await websocket.send_json({"type": "ready"})

    buffer = bytearray()
    mimetype: str = "audio/webm"
    last_interim_at = 0.0
    last_sent_text = ""
    transcribe_lock = asyncio.Lock()
    interim_task: asyncio.Task[None] | None = None
    closed = False

    def _suffix_for(mt: str) -> str:
        if "mp4" in mt:
            return ".m4a"
        if "ogg" in mt:
            return ".ogg"
        if "wav" in mt:
            return ".wav"
        return ".webm"

    async def _run_transcribe(final: bool) -> None:
        """Run faster-whisper on the current buffer and send a message.

        Serialised by ``transcribe_lock``. Safe to call from multiple
        call sites — only one actually runs at a time; queued callers
        are quietly dropped unless ``final`` is True.
        """
        nonlocal last_sent_text
        if not buffer:
            if final:
                with suppress_ws_errors():
                    await websocket.send_json({"type": "final", "text": ""})
            return
        # Non-final: skip if another tick is already running. Final: wait.
        if not final and transcribe_lock.locked():
            return
        async with transcribe_lock:
            snapshot = bytes(buffer)
            suffix = _suffix_for(mimetype)
            try:
                text = await asyncio.to_thread(transcribe_buffer_sync, snapshot, suffix)
            except Exception as exc:
                log.warning("Streaming transcribe tick failed: %s", exc)
                return
            text = (text or "").strip()
            if final:
                last_sent_text = text
                with suppress_ws_errors():
                    await websocket.send_json({"type": "final", "text": text})
                return
            # Only emit an interim update if the text actually changed —
            # avoids jitter from the model re-transcribing the same audio.
            if text and text != last_sent_text:
                last_sent_text = text
                with suppress_ws_errors():
                    await websocket.send_json({"type": "interim", "text": text})

    def _maybe_schedule_interim() -> None:
        """Kick off an interim transcription tick if enough time has elapsed."""
        nonlocal last_interim_at, interim_task
        if closed:
            return
        now = time.monotonic()
        if now - last_interim_at < _INTERIM_INTERVAL:
            return
        if interim_task and not interim_task.done():
            return
        last_interim_at = now
        interim_task = asyncio.create_task(_run_transcribe(final=False))

    try:
        while True:
            msg = await websocket.receive()
            # Starlette delivers disconnect messages as a type, not an exception
            if msg.get("type") == "websocket.disconnect":
                break

            if "bytes" in msg and msg["bytes"] is not None:
                chunk: bytes = msg["bytes"]
                if len(buffer) + len(chunk) > _MAX_BUFFER_BYTES:
                    await websocket.send_json({
                        "type": "error",
                        "detail": "Audio buffer exceeded 40 MB limit.",
                    })
                    break
                buffer.extend(chunk)
                _maybe_schedule_interim()
                continue

            if "text" in msg and msg["text"] is not None:
                try:
                    payload = json.loads(msg["text"])
                except json.JSONDecodeError:
                    continue
                ptype = payload.get("type")
                if ptype == "mimetype":
                    candidate = payload.get("value")
                    if isinstance(candidate, str) and candidate:
                        mimetype = candidate
                elif ptype == "cancel":
                    buffer.clear()
                    with suppress_ws_errors():
                        await websocket.send_json({"type": "final", "text": ""})
                    break
                elif ptype == "done":
                    # Wait out any in-flight interim task then run a final pass
                    if interim_task and not interim_task.done():
                        with contextlib.suppress(Exception):
                            await interim_task
                    await _run_transcribe(final=True)
                    break
    except WebSocketDisconnect:
        pass
    except Exception:
        log.exception("Transcribe WS handler crashed")
        with suppress_ws_errors():
            await websocket.send_json({"type": "error", "detail": "Internal error"})
    finally:
        closed = True
        if interim_task and not interim_task.done():
            interim_task.cancel()
            with contextlib.suppress(Exception):
                await interim_task
        with suppress_ws_errors():
            await websocket.close()
