"""Opt-in, content-free send diagnostics. No wire/protocol or pacing changes."""
from __future__ import annotations

import asyncio
import json
import logging
import os
import time
from dataclasses import dataclass

log = logging.getLogger('uvicorn.error.device_gateway.speech_trace')


@dataclass
class Stream:
    speech_id: str
    device_id: str
    source: str
    origin: float
    chunks: int = 0
    raw_bytes: int = 0
    max_send_ms: float = 0
    max_lock_ms: float = 0
    max_drift_ms: float = 0
    end_sent: bool = False


class SpeechTrace:
    def __init__(self, enabled=False, *, clock=time.monotonic, sink=None, every=8):
        self.enabled, self.clock, self.sink = enabled, clock, sink
        self.every = max(1, every)
        self.streams = {}

    @classmethod
    def from_env(cls):
        # Full JSONL can be captured from journal using EVERY=1; no payload file.
        return cls(os.getenv('GATEWAY_SPEECH_TRACE') == '1',
                   every=1 if os.getenv('GATEWAY_SPEECH_TRACE_EVERY') == '1' else 8)

    def begin(self, device, speech_id, source):
        if self.enabled:
            self.streams[device.connection_id] = Stream(speech_id, device.device_id, source, self.clock())
            self.emit(dict(event='begin', speech_id=speech_id, device_id=device.device_id, source=source))

    def emit(self, record, important=True):
        if self.sink is not None:
            self.sink(dict(record))  # Test sink receives ALL records, including unsampled sends.
        if important:
            log.info('[SPEECH_TRACE] %s', json.dumps(record, separators=(',', ':'), ensure_ascii=True))

    def finish(self, device, speech_id):
        stream = self.streams.get(device.connection_id)
        if stream is None or stream.speech_id != speech_id:
            return
        self.streams.pop(device.connection_id, None)
        self.emit(dict(event='summary', speech_id=stream.speech_id, device_id=stream.device_id,
                       source=stream.source, completed_chunks=stream.chunks,
                       raw_pcm_bytes=stream.raw_bytes, max_websocket_send_ms=stream.max_send_ms,
                       max_device_lock_wait_ms=stream.max_lock_ms, max_drift_ms=stream.max_drift_ms,
                       speech_end_sent=stream.end_sent, elapsed_ms=(self.clock()-stream.origin)*1000))

    async def send(self, device, payload, *, scheduled_elapsed_ms=None, raw_pcm_bytes=0):
        stream = self.streams.get(device.connection_id) if self.enabled else None
        if stream is None:
            async with device.send_lock:
                await asyncio.wait_for(device.websocket.send_json(payload), 3)
            return
        # Serialize with exactly Starlette's settings for UTF-8 application bytes.
        # This is not compressed WebSocket/TLS/TCP wire size. Never retain the value.
        wire_bytes = len(json.dumps(payload, separators=(',', ':'), ensure_ascii=False).encode())
        attempt = self.clock()
        acquired = complete = None
        error = None
        kind = payload.get('type')
        sequence = payload.get('sequence') if kind == 'speech_chunk' else None
        is_stream = payload.get('speech_id') == stream.speech_id
        if not is_stream:
            scheduled_elapsed_ms, raw_pcm_bytes = None, 0
        try:
            async with device.send_lock:
                acquired = self.clock()
                await asyncio.wait_for(device.websocket.send_json(payload), 3)
                complete = self.clock()
        except BaseException as exc:
            error = type(exc).__name__  # Never exception text (may contain payload or credentials).
            raise
        finally:
            finished = self.clock() if complete is None else complete
            lock_ms = ((finished if acquired is None else acquired)-attempt)*1000
            send_ms = None if acquired is None else (finished-acquired)*1000
            drift = None if scheduled_elapsed_ms is None else (attempt-stream.origin)*1000-scheduled_elapsed_ms
            stream.max_send_ms = max(stream.max_send_ms, send_ms or 0)
            stream.max_lock_ms = max(stream.max_lock_ms, lock_ms)
            stream.max_drift_ms = max(stream.max_drift_ms, drift or 0)
            if complete is not None and is_stream:
                stream.chunks += kind == 'speech_chunk'
                stream.raw_bytes += raw_pcm_bytes
                stream.end_sent |= kind == 'speech_end'
            record = dict(event='send', speech_id=stream.speech_id, device_id=stream.device_id,
                          source=stream.source, message_type=kind, sequence=sequence,
                          chunks=payload.get('chunks') if kind == 'speech_end' else None,
                          raw_pcm_bytes=raw_pcm_bytes, wire_payload_bytes=wire_bytes,
                          scheduled_elapsed_ms=scheduled_elapsed_ms,
                          send_attempt_ms=(attempt-stream.origin)*1000,
                          device_lock_wait_ms=lock_ms, websocket_send_ms=send_ms,
                          send_complete_ms=None if complete is None else (complete-stream.origin)*1000,
                          drift_from_nominal_ms=drift, error=error)
            important = (kind != 'speech_chunk' or sequence is None or sequence % self.every == 0
                         or error is not None or (send_ms or 0) > 150 or lock_ms > 100 or (drift or 0) > 150)
            self.emit(record, important)


trace = SpeechTrace.from_env()