"""Server-side streaming session for /v1/realtime.

Drives the shared LocalAgreementCommitter over PCM16 audio arriving as binary
WebSocket frames, emitting delta/completed events whose names mirror the OpenAI
Realtime transcription schema so existing tooling interoperates.

The incoming audio is kept as a list of chunks with a rolling tail (see
_evict_committed): chunks entirely before min(committer anchor, tail window)
are dropped, so `base` tracks how many samples were evicted off the front and
the kept buffer stays bounded regardless of session length. The session is
also hard-capped at MAX_SESSION_SECONDS.
"""
from __future__ import annotations

import asyncio
import functools

import numpy as np
from loguru import logger

from services.local_agreement import (
    MIN_NEW_AUDIO_S,
    POLL_INTERVAL,
    SAMPLE_RATE,
    LocalAgreementCommitter,
)

# Client → server event types (mirrors OpenAI Realtime input_audio_buffer.*)
EVT_APPEND = "input_audio_buffer.append"
EVT_COMMIT = "input_audio_buffer.commit"   # finalize (client stopped speaking)
EVT_CLEAR = "input_audio_buffer.clear"

# Server → client event types
EVT_DELTA = "conversation.item.input_audio_transcription.delta"
EVT_DONE = "conversation.item.input_audio_transcription.completed"
EVT_ERROR = "error"


# Bound a single streaming session so one client cannot grow memory without
# limit or hold resources forever. 10 min of 16 kHz mono is ~19 MB of float32.
MAX_SESSION_SECONDS = 600
# Keep only a rolling tail in the WS buffer — the committer never needs more
# than its own hard cap, and evicting committed audio keeps concat cheap and
# memory flat. `base` tracks how many samples were dropped off the front.
_ROLLING_KEEP_SECONDS = 40


class RealtimeSession:
    """One WebSocket streaming session. Frame ingestion and the periodic
    inference loop run concurrently. `_buf_lock` guards the incoming-frame
    buffer; `_infer_lock` serializes ALL committer access (the committer is
    documented not-thread-safe, and two inferences must never overlap)."""

    def __init__(self, engine, language: str | None) -> None:
        self._engine = engine
        self._committer = LocalAgreementCommitter(language)
        # Incoming PCM as a list of chunks (append is O(1); we concat only the
        # kept tail on demand — no O(n^2) whole-buffer copy per frame).
        self._chunks: list[np.ndarray] = []
        self._chunk_samples = 0
        self._base = 0                 # samples evicted off the front (absolute)
        self._buf_lock = asyncio.Lock()
        self._infer_lock = asyncio.Lock()
        self._closed = False
        self._committed_text: list[str] = []

    def add_pcm16(self, raw: bytes) -> None:
        """Append little-endian int16 PCM (16 kHz mono) as float32 [-1, 1]."""
        if not raw:
            return
        pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
        self._chunks.append(pcm)
        self._chunk_samples += len(pcm)
        self._evict_committed()

    def _evict_committed(self) -> None:
        """Drop chunks entirely before the eviction cutoff so the kept tail
        stays bounded, but NEVER past the committer's anchor (audio at/after the
        anchor has not been processed yet — evicting it would silently lose
        speech). The cutoff is therefore min(anchor, rolling-tail-start): we
        evict as much as the rolling window allows, but never beyond what the
        committer still needs, even when inference lags far behind under load."""
        tail_start = (self._base + self._chunk_samples
                      - _ROLLING_KEEP_SECONDS * SAMPLE_RATE)
        keep_from = min(self._committer.anchor_samples, tail_start)
        while self._chunks and self._base + len(self._chunks[0]) <= keep_from:
            dropped = self._chunks.pop(0)
            self._base += len(dropped)
            self._chunk_samples -= len(dropped)

    def _snapshot(self) -> tuple[np.ndarray, int]:
        if not self._chunks:
            return np.zeros(0, dtype=np.float32), self._base
        return np.concatenate(self._chunks), self._base

    def _captured(self) -> int:
        return self._base + self._chunk_samples

    async def _infer(self, *, final: bool) -> str:
        loop = asyncio.get_running_loop()
        # Serialize the whole inference (snapshot + committer.process): the
        # committer is single-thread-only, and admission is acquired here so an
        # idle live session never holds the global slot.
        async with self._infer_lock:
            async with self._buf_lock:
                full, base = self._snapshot()
            async with self._engine.admit():
                delta = await loop.run_in_executor(
                    None,
                    functools.partial(
                        self._committer.process,
                        full, base, self._engine.service.transcribe_window_sync,
                        final=final,
                    ),
                )
        if delta:
            self._committed_text.append(delta)
        return delta

    async def run(self, websocket) -> None:
        """Main session loop: consume frames and periodically infer."""
        infer_task = asyncio.create_task(self._infer_loop(websocket))
        try:
            while not self._closed:
                message = await websocket.receive()
                if message.get("type") == "websocket.disconnect":
                    break
                if (data := message.get("bytes")) is not None:
                    async with self._buf_lock:
                        self.add_pcm16(data)
                    if self._captured() > MAX_SESSION_SECONDS * SAMPLE_RATE:
                        # Session too long: finalize and stop instead of growing.
                        await self._finalize(websocket)
                        break
                elif (text := message.get("text")) is not None:
                    await self._handle_control(websocket, text)
        finally:
            self._closed = True
            infer_task.cancel()
            try:
                await infer_task
            except asyncio.CancelledError:
                pass

    async def _finalize(self, websocket) -> None:
        try:
            delta = await self._infer(final=True)
        except Exception as e:  # noqa: BLE001 — a busy/errored engine must tell
            logger.warning("Realtime finalize failed: {}", e)
            await self._send(websocket, EVT_ERROR,
                             {"error": {"message": "Could not finalize transcription"}})
            self._closed = True
            return
        if delta:
            await self._send(websocket, EVT_DELTA, {"delta": delta})
        await self._send(websocket, EVT_DONE,
                         {"transcript": "".join(self._committed_text).strip()})
        self._closed = True

    async def _handle_control(self, websocket, text: str) -> None:
        import json
        try:
            evt = json.loads(text)
        except ValueError:
            return
        etype = evt.get("type")
        if etype == EVT_COMMIT:
            await self._finalize(websocket)
        elif etype == EVT_CLEAR:
            async with self._infer_lock:
                async with self._buf_lock:
                    self._chunks = []
                    self._chunk_samples = 0
                    self._base = 0
                    self._committer.reset()
                    self._committed_text = []

    async def _infer_loop(self, websocket) -> None:
        """Periodic incremental inference while frames stream in."""
        await asyncio.sleep(POLL_INTERVAL)
        min_new = int(MIN_NEW_AUDIO_S * SAMPLE_RATE)
        while not self._closed:
            if self._captured() - self._committer.last_infer_abs >= min_new:
                try:
                    delta = await self._infer(final=False)
                    if delta:
                        await self._send(websocket, EVT_DELTA, {"delta": delta})
                except Exception as e:  # noqa: BLE001
                    logger.warning("Realtime inference failed: {}", e)
            await asyncio.sleep(POLL_INTERVAL)

    async def _send(self, websocket, event_type: str, payload: dict) -> None:
        import json
        try:
            await websocket.send_text(json.dumps({"type": event_type, **payload}))
        except Exception:  # noqa: BLE001 — client may have vanished
            self._closed = True
