"""Capture raw xAI TTS WebSocket events as JSONL.

This collector records direct client/provider traffic. It does not use an audio
emitter, retries, or a gap timeout; each JSONL line is an event for one call.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import binascii
import hashlib
import json
import math
import os
import sys
import time
import uuid
from collections import Counter
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.parse import urlencode

import aiohttp
from dotenv import load_dotenv

SAMPLE_RATE = 24_000
PCM_BYTES_PER_SAMPLE = 2
DEFAULT_TEXT = "This is a raw capture of TTS WebSocket events."


@dataclass(frozen=True, slots=True)
class Config:
    api_key: str
    websocket_url: str
    voice: str
    language: str
    text: str
    calls: int
    concurrency: int
    prewarm: bool
    event_timeout_s: float | None
    connect_timeout_s: float | None
    output: Path
    summary_output: Path
    include_delta_base64: bool
    initial_playout_buffer_ms: float


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds")


def elapsed_ms(started_at: float) -> int:
    return max(0, round((time.perf_counter() - started_at) * 1000))


def bool_query(value: bool) -> str:
    return "true" if value else "false"


def build_url(config: Config) -> str:
    query = {
        "voice": config.voice,
        "language": config.language,
        "codec": "pcm",
        "sample_rate": str(SAMPLE_RATE),
        "optimize_streaming_latency": "1",
        "text_normalization": bool_query(True),
    }
    return f"{config.websocket_url}?{urlencode(query)}"


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(
        description="Capture raw xAI TTS WebSocket events as JSONL."
    )
    result.add_argument("--env-file", default=".env.dev")
    result.add_argument("--api-key", default="")
    result.add_argument("--websocket-url", default="")
    result.add_argument("--voice", default="")
    result.add_argument("--language", default="")
    result.add_argument("--text", default=DEFAULT_TEXT)
    result.add_argument("--calls", type=int, default=1)
    result.add_argument("--concurrency", type=int, default=1)
    result.add_argument(
        "--prewarm",
        action="store_true",
        help="Open every WebSocket before the measured window. Requires --calls to equal --concurrency.",
    )
    result.add_argument(
        "--event-timeout-ms",
        type=int,
        default=65_000,
        help="Safety limit for each ws.receive call; 0 disables it. This is not a gap timeout.",
    )
    result.add_argument(
        "--connect-timeout-ms",
        type=int,
        default=15_000,
        help="Safety limit for each WebSocket handshake; 0 disables it.",
    )
    result.add_argument("--output", default=".run/xai-tts-ws-capture/events.jsonl")
    result.add_argument(
        "--summary-output",
        default="",
        help="JSON summary output. Default: <output>.summary.json.",
    )
    result.add_argument(
        "--initial-playout-buffer-ms",
        type=float,
        default=0.0,
        help="Optional intentional prebuffer for the underflow simulation (default: 0; playout starts at the first non-empty audio.delta).",
    )
    result.add_argument(
        "--include-delta-base64",
        action="store_true",
        help="Include the full Base64 delta in each audio.delta event. This can produce very large files.",
    )
    return result


def load_config(args: argparse.Namespace) -> Config:
    env_file = Path(args.env_file)
    if env_file.exists():
        load_dotenv(env_file, override=False)
    api_key = str(args.api_key or os.getenv("XAI_API_KEY", "")).strip()
    websocket_url = str(
        args.websocket_url or os.getenv("XAI_WEBSOCKET_URL", "")
    ).strip()
    if not api_key:
        raise RuntimeError("XAI_API_KEY is missing; provide --api-key or --env-file.")
    if not websocket_url:
        raise RuntimeError(
            "XAI_WEBSOCKET_URL is missing; provide --websocket-url or --env-file."
        )
    if not str(args.text).strip():
        raise ValueError("--text must not be empty.")
    if int(args.calls) < 1:
        raise ValueError("--calls must be at least 1.")
    if int(args.concurrency) < 1:
        raise ValueError("--concurrency must be at least 1.")
    timeout_ms = max(0, int(args.event_timeout_ms))
    connect_timeout_ms = max(0, int(args.connect_timeout_ms))
    initial_playout_buffer_ms = float(args.initial_playout_buffer_ms)
    if not math.isfinite(initial_playout_buffer_ms) or initial_playout_buffer_ms < 0:
        raise ValueError(
            "--initial-playout-buffer-ms must be a finite value greater than or equal to 0."
        )
    output = Path(args.output)
    summary_output = (
        Path(args.summary_output)
        if args.summary_output
        else output.with_suffix(".summary.json")
    )
    if output.absolute() == summary_output.absolute():
        raise ValueError("--output and --summary-output must refer to different files.")
    return Config(
        api_key=api_key,
        websocket_url=websocket_url,
        voice=str(args.voice or os.getenv("XAI_TTS_VOICE", "ara")).strip() or "ara",
        language=str(args.language or os.getenv("XAI_TTS_LANGUAGE", "pt-BR")).strip()
        or "pt-BR",
        text=str(args.text),
        calls=int(args.calls),
        concurrency=int(args.concurrency),
        prewarm=bool(args.prewarm),
        event_timeout_s=(timeout_ms / 1000) if timeout_ms else None,
        connect_timeout_s=(connect_timeout_ms / 1000) if connect_timeout_ms else None,
        output=output,
        summary_output=summary_output,
        include_delta_base64=bool(args.include_delta_base64),
        initial_playout_buffer_ms=initial_playout_buffer_ms,
    )


class JsonlWriter:
    def __init__(self, output: Path) -> None:
        output.parent.mkdir(parents=True, exist_ok=True)
        self._handle = output.open("w", encoding="utf-8")
        self._lock = asyncio.Lock()

    async def write(self, record: dict[str, Any]) -> None:
        async with self._lock:
            self._handle.write(
                json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n"
            )
            self._handle.flush()

    def close(self) -> None:
        self._handle.close()


async def receive_message(
    ws: aiohttp.ClientWebSocketResponse, timeout_s: float | None
) -> aiohttp.WSMessage:
    if timeout_s is None:
        return await ws.receive()
    return await asyncio.wait_for(ws.receive(), timeout=timeout_s)


def event_base(
    *, run_id: str, call_id: str, call_number: int, sequence: int, started_at: float
) -> dict[str, Any]:
    return {
        "run_id": run_id,
        "call_id": call_id,
        "call_number": call_number,
        "sequence": sequence,
        "received_at_utc": utc_now(),
        "received_after_call_start_ms": elapsed_ms(started_at),
    }


async def capture_call(
    *,
    session: aiohttp.ClientSession,
    config: Config,
    writer: JsonlWriter,
    run_id: str,
    call_number: int,
    call_id: str,
    limiter: asyncio.Semaphore,
    prewarmed_ws: aiohttp.ClientWebSocketResponse | None = None,
) -> str:
    async with limiter:
        started_at = time.perf_counter()
        sequence = 0
        status = "unknown"
        ws = prewarmed_ws

        async def write(kind: str, **fields: Any) -> None:
            nonlocal sequence
            sequence += 1
            record = event_base(
                run_id=run_id,
                call_id=call_id,
                call_number=call_number,
                sequence=sequence,
                started_at=started_at,
            )
            record["event"] = kind
            record.update(fields)
            await writer.write(record)

        try:
            if ws is None:
                await write("client.connecting", websocket_url=config.websocket_url)
                connect_started_at = time.perf_counter()
                ws = await session.ws_connect(
                    build_url(config),
                    headers={"Authorization": f"Bearer {config.api_key}"},
                    heartbeat=20,
                )
                await write(
                    "client.connected", connect_ms=elapsed_ms(connect_started_at)
                )
            else:
                await write("client.using_prewarmed_connection")

            await ws.send_json({"type": "text.clear"})
            await write("client.sent", message_type="text.clear")
            clear_acknowledged = False
            while not clear_acknowledged:
                message = await receive_message(ws, config.event_timeout_s)
                now = time.perf_counter()
                if message.type != aiohttp.WSMsgType.TEXT:
                    await write(
                        "transport.message",
                        ws_type=message.type.name,
                        data=str(message.data),
                    )
                    raise RuntimeError(
                        f"Expected audio.clear, received WebSocket {message.type.name}."
                    )
                try:
                    payload = json.loads(str(message.data))
                except json.JSONDecodeError:
                    await write("provider.invalid_json", raw_text=str(message.data))
                    raise
                await write(
                    "provider.event",
                    message_type=str(payload.get("type") or ""),
                    provider_message=payload,
                    received_monotonic_ms=round(now * 1000),
                )
                clear_acknowledged = payload.get("type") == "audio.clear"
                if payload.get("type") == "error":
                    raise RuntimeError(str(payload.get("message") or payload))

            await ws.send_json({"type": "text.delta", "delta": config.text})
            await write(
                "client.sent", message_type="text.delta", text_chars=len(config.text)
            )
            await ws.send_json({"type": "text.done"})
            await write("client.sent", message_type="text.done")

            while True:
                try:
                    message = await receive_message(ws, config.event_timeout_s)
                except asyncio.TimeoutError:
                    await write(
                        "client.receive_timeout",
                        timeout_ms=None
                        if config.event_timeout_s is None
                        else round(config.event_timeout_s * 1000),
                    )
                    status = "receive_timeout"
                    break

                now = time.perf_counter()
                if message.type != aiohttp.WSMsgType.TEXT:
                    await write(
                        "transport.message",
                        ws_type=message.type.name,
                        data=str(message.data),
                    )
                    status = f"transport_{message.type.name.lower()}"
                    break

                try:
                    payload = json.loads(str(message.data))
                except json.JSONDecodeError:
                    await write("provider.invalid_json", raw_text=str(message.data))
                    status = "invalid_json"
                    break

                message_type = str(payload.get("type") or "")
                if message_type == "audio.delta":
                    encoded = str(payload.get("delta") or "")
                    try:
                        pcm = base64.b64decode(encoded, validate=True)
                    except (binascii.Error, ValueError) as exc:
                        await write(
                            "provider.audio_delta_invalid",
                            message_type=message_type,
                            base64_characters=len(encoded),
                            error=str(exc),
                            provider_message={
                                key: value
                                for key, value in payload.items()
                                if key != "delta"
                            },
                        )
                        status = "invalid_audio_delta"
                        break
                    record: dict[str, Any] = {
                        "message_type": message_type,
                        "base64_characters": len(encoded),
                        "audio_bytes": len(pcm),
                        "audio_ms": round(
                            len(pcm) * 1000 / (SAMPLE_RATE * PCM_BYTES_PER_SAMPLE), 3
                        ),
                        "empty_audio": not bool(pcm),
                        "audio_sha256": hashlib.sha256(pcm).hexdigest(),
                        "provider_message": {
                            key: value
                            for key, value in payload.items()
                            if key != "delta"
                        },
                        "received_monotonic_ms": round(now * 1000),
                    }
                    if config.include_delta_base64:
                        record["delta_base64"] = encoded
                    await write("provider.audio_delta", **record)
                    continue

                await write(
                    "provider.event",
                    message_type=message_type,
                    provider_message=payload,
                    received_monotonic_ms=round(now * 1000),
                )
                if message_type == "audio.done":
                    status = "audio_done"
                    break
                if message_type == "error":
                    status = "provider_error"
                    break
        except Exception as exc:
            if status == "unknown":
                status = "exception"
            await write(
                "client.exception", error_type=type(exc).__name__, error=str(exc)
            )
        finally:
            if ws is not None and not ws.closed:
                await ws.close()
                await write("client.closed")
            await write(
                "client.finished", status=status, elapsed_ms=elapsed_ms(started_at)
            )
        return status


async def prewarm_connection(
    *,
    session: aiohttp.ClientSession,
    config: Config,
    writer: JsonlWriter,
    run_id: str,
    call_number: int,
    call_id: str,
) -> aiohttp.ClientWebSocketResponse:
    started_at = time.perf_counter()
    await writer.write(
        {
            **event_base(
                run_id=run_id,
                call_id=call_id,
                call_number=call_number,
                sequence=-1,
                started_at=started_at,
            ),
            "event": "client.prewarm_connecting",
            "websocket_url": config.websocket_url,
        }
    )
    try:
        ws = await session.ws_connect(
            build_url(config),
            headers={"Authorization": f"Bearer {config.api_key}"},
            heartbeat=20,
        )
    except Exception as exc:
        await writer.write(
            {
                **event_base(
                    run_id=run_id,
                    call_id=call_id,
                    call_number=call_number,
                    sequence=0,
                    started_at=started_at,
                ),
                "event": "client.prewarm_failed",
                "error_type": type(exc).__name__,
                "error": str(exc),
            }
        )
        raise
    await writer.write(
        {
            **event_base(
                run_id=run_id,
                call_id=call_id,
                call_number=call_number,
                sequence=0,
                started_at=started_at,
            ),
            "event": "client.prewarm_ready",
            "connect_ms": elapsed_ms(started_at),
        }
    )
    return ws


def percentile(values: list[float], percentile_value: float) -> float | None:
    """Return the nearest-rank percentile without introducing a dependency."""
    if not values:
        return None
    ordered = sorted(values)
    index = round((len(ordered) - 1) * percentile_value)
    return ordered[max(0, min(index, len(ordered) - 1))]


def round_ms(value: float | None) -> float | None:
    return None if value is None else round(value, 3)


def replay_underflow(
    audio_frames: list[tuple[float, float]], initial_buffer_ms: float
) -> tuple[int, float, float]:
    """Replay raw arrivals; with a zero prebuffer, playout starts at first audio.delta."""
    buffered_ms = 0.0
    playout_started = False
    previous_arrival_ms: float | None = None
    events = 0
    total_underflow_ms = 0.0
    largest_underflow_ms = 0.0

    for arrived_at_ms, audio_ms in audio_frames:
        if playout_started:
            assert previous_arrival_ms is not None
            elapsed_ms = arrived_at_ms - previous_arrival_ms
            missing_audio_ms = max(0.0, elapsed_ms - buffered_ms)
            if missing_audio_ms:
                events += 1
                total_underflow_ms += missing_audio_ms
                largest_underflow_ms = max(largest_underflow_ms, missing_audio_ms)
            buffered_ms = max(0.0, buffered_ms - elapsed_ms) + audio_ms
            previous_arrival_ms = arrived_at_ms
            continue

        buffered_ms += audio_ms
        if buffered_ms >= initial_buffer_ms:
            playout_started = True
            previous_arrival_ms = arrived_at_ms

    return events, total_underflow_ms, largest_underflow_ms


def metric_distribution(values: list[float]) -> dict[str, float | int | None]:
    return {
        "count": len(values),
        "p50": round_ms(percentile(values, 0.50)),
        "p95": round_ms(percentile(values, 0.95)),
        "max": round_ms(max(values)) if values else None,
    }


def read_finite_float(record: dict[str, Any], field: str) -> float | None:
    """Read a finite numeric field without allowing a malformed capture to break analysis."""
    try:
        value = float(record[field])
    except (KeyError, TypeError, ValueError):
        return None
    return value if math.isfinite(value) else None


def event_sequence(record: dict[str, Any]) -> int:
    """Sort malformed or missing sequence values before regular events."""
    try:
        return int(record.get("sequence", -2))
    except (TypeError, ValueError):
        return -2


def display_metric(value: float | None, unit: str = "") -> str:
    return "n/a" if value is None else f"{value}{unit}"


def summarize_capture(config: Config, run_id: str) -> dict[str, Any]:
    """Build a provider-facing summary exclusively from the raw JSONL capture."""
    calls: dict[str, list[tuple[int, dict[str, Any]]]] = {}
    run_events: list[dict[str, Any]] = []
    analysis_warnings: list[str] = []
    with config.output.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError as exc:
                analysis_warnings.append(
                    f"Skipped invalid JSONL record at line {line_number}: {exc.msg}."
                )
                continue
            if not isinstance(record, dict):
                analysis_warnings.append(
                    f"Skipped non-object JSONL record at line {line_number}."
                )
                continue
            call_id = record.get("call_id")
            if call_id:
                calls.setdefault(str(call_id), []).append((line_number, record))
            else:
                run_events.append(record)

    status_counts: Counter[str] = Counter()
    error_counts: Counter[str] = Counter()
    ttfb_ms: list[float] = []
    rtf: list[float] = []
    audio_duration_ms: list[float] = []
    frame_duration_ms: list[float] = []
    frame_gaps_ms: list[float] = []
    per_call_max_gaps_ms: list[float] = []
    largest_gap_observation: dict[str, Any] | None = None
    frames_per_call: list[float] = []
    calls_with_audio = 0
    calls_with_underflow = 0
    underflow_events = 0
    total_underflow_ms = 0.0
    largest_underflow_ms = 0.0
    empty_audio_deltas = 0

    for call_id, call_records in calls.items():
        call_records.sort(key=lambda item: event_sequence(item[1]))
        records = [record for _, record in call_records]
        for record in records:
            event = str(record.get("event") or "")
            message_type = str(record.get("message_type") or "")
            if event == "client.finished":
                status_counts[str(record.get("status") or "unknown")] += 1
            if event == "client.prewarm_failed":
                error_counts["prewarm_failed"] += 1
            elif event == "client.exception":
                error_counts["client_exception"] += 1
            elif event == "client.receive_timeout":
                error_counts["receive_timeout"] += 1
            elif event == "provider.audio_delta_invalid":
                error_counts["invalid_audio_delta"] += 1
            elif event == "provider.invalid_json":
                error_counts["invalid_json"] += 1
            elif event == "transport.message":
                error_counts[
                    f"transport_{str(record.get('ws_type') or 'unknown').lower()}"
                ] += 1
            elif event == "provider.event" and message_type == "error":
                error_counts["provider_error"] += 1
            elif event == "provider.audio_delta" and record.get("empty_audio"):
                empty_audio_deltas += 1

        text_done_at_ms = None
        for line_number, record in call_records:
            if (
                record.get("event") != "client.sent"
                or record.get("message_type") != "text.done"
            ):
                continue
            text_done_at_ms = read_finite_float(record, "received_after_call_start_ms")
            if text_done_at_ms is None:
                analysis_warnings.append(
                    f"Skipped invalid text.done timestamp at line {line_number} for {call_id}."
                )
            break
        audio_frame_records: list[dict[str, Any]] = []
        for line_number, record in call_records:
            if record.get("event") != "provider.audio_delta":
                continue
            audio_bytes = read_finite_float(record, "audio_bytes")
            received_at_ms = read_finite_float(record, "received_after_call_start_ms")
            frame_audio_ms = read_finite_float(record, "audio_ms")
            if audio_bytes is None or received_at_ms is None or frame_audio_ms is None:
                analysis_warnings.append(
                    f"Skipped malformed audio.delta at line {line_number} for {call_id}."
                )
                continue
            if audio_bytes > 0 and frame_audio_ms >= 0:
                audio_frame_records.append(record)
        audio_frames = [
            (float(record["received_after_call_start_ms"]), float(record["audio_ms"]))
            for record in audio_frame_records
        ]
        if text_done_at_ms is None or not audio_frames:
            continue

        calls_with_audio += 1
        frames_per_call.append(float(len(audio_frames)))
        duration_ms = sum(frame[1] for frame in audio_frames)
        audio_duration_ms.append(duration_ms)
        frame_duration_ms.extend(frame[1] for frame in audio_frames)
        ttfb_ms.append(audio_frames[0][0] - text_done_at_ms)
        if duration_ms:
            rtf.append((audio_frames[-1][0] - text_done_at_ms) / duration_ms)
        gaps = [
            audio_frames[index][0] - audio_frames[index - 1][0]
            for index in range(1, len(audio_frames))
        ]
        frame_gaps_ms.extend(gaps)
        if gaps:
            per_call_max_gaps_ms.append(max(gaps))
            largest_gap_index = max(
                range(1, len(audio_frames)),
                key=lambda index: audio_frames[index][0] - audio_frames[index - 1][0],
            )
            largest_gap_ms = (
                audio_frames[largest_gap_index][0]
                - audio_frames[largest_gap_index - 1][0]
            )
            if (
                largest_gap_observation is None
                or largest_gap_ms > largest_gap_observation["gap_ms"]
            ):
                previous_record = audio_frame_records[largest_gap_index - 1]
                current_record = audio_frame_records[largest_gap_index]
                largest_gap_observation = {
                    "gap_ms": round_ms(largest_gap_ms),
                    "call_id": call_id,
                    "previous_audio_delta_sequence": previous_record.get("sequence"),
                    "previous_arrived_after_call_start_ms": audio_frames[
                        largest_gap_index - 1
                    ][0],
                    "previous_frame_audio_ms": audio_frames[largest_gap_index - 1][1],
                    "next_audio_delta_sequence": current_record.get("sequence"),
                    "next_arrived_after_call_start_ms": audio_frames[largest_gap_index][
                        0
                    ],
                    "next_frame_audio_ms": audio_frames[largest_gap_index][1],
                }
        events, total_ms, largest_ms = replay_underflow(
            audio_frames, config.initial_playout_buffer_ms
        )
        if events:
            calls_with_underflow += 1
            underflow_events += events
            total_underflow_ms += total_ms
            largest_underflow_ms = max(largest_underflow_ms, largest_ms)

    run_failed = next(
        (record for record in run_events if record.get("event") == "run.failed"), None
    )
    completed_calls = status_counts.get("audio_done", 0)
    observed_errors = sum(error_counts.values())
    return {
        "summary_version": 1,
        "language": "en",
        "source_events_jsonl": str(config.output),
        "run_id": run_id,
        "configuration": {
            "requested_calls": config.calls,
            "concurrency": config.concurrency,
            "prewarm": config.prewarm,
            "initial_playout_buffer_ms": config.initial_playout_buffer_ms,
            "event_timeout_ms": None
            if config.event_timeout_s is None
            else round(config.event_timeout_s * 1000),
            "connect_timeout_ms": None
            if config.connect_timeout_s is None
            else round(config.connect_timeout_s * 1000),
            "pcm_format": f"{SAMPLE_RATE} Hz, mono, signed 16-bit PCM",
        },
        "completion": {
            "calls_observed": len(calls),
            "calls_with_audio": calls_with_audio,
            "audio_done": completed_calls,
            "audio_done_rate": round(completed_calls / config.calls, 4)
            if config.calls
            else None,
            "finished_statuses": dict(sorted(status_counts.items())),
            "run_failed": bool(run_failed),
            "run_failure": None if run_failed is None else run_failed.get("error"),
        },
        "latency": {
            "ttfb_ms": metric_distribution(ttfb_ms),
            "rtf": metric_distribution(rtf),
            "generated_audio_ms_per_call": metric_distribution(audio_duration_ms),
        },
        "frame_delivery": {
            "raw_interframe_arrival_gap_definition": "Timestamp difference between consecutive non-empty provider.audio_delta messages for the same call. It excludes TTFB and does not include either frame's audio duration.",
            "nonempty_audio_frames_per_call": metric_distribution(frames_per_call),
            "audio_ms_per_frame": metric_distribution(frame_duration_ms),
            "raw_interframe_arrival_gap_ms": metric_distribution(frame_gaps_ms),
            "per_call_max_raw_interframe_arrival_gap_ms": metric_distribution(
                per_call_max_gaps_ms
            ),
            "largest_raw_interframe_arrival_gap": largest_gap_observation,
            "empty_audio_delta_count": empty_audio_deltas,
        },
        "simulated_playout": {
            "definition": "TTFB is excluded. The replay clock starts when the first non-empty audio.delta arrives. If initial_buffer_ms is greater than zero, playout is intentionally delayed until that much PCM has accumulated; otherwise it starts immediately at the first frame.",
            "initial_buffer_ms": config.initial_playout_buffer_ms,
            "calls_with_underflow": calls_with_underflow,
            "calls_with_underflow_rate": round(calls_with_underflow / config.calls, 4)
            if config.calls
            else None,
            "underflow_events": underflow_events,
            "total_underflow_ms": round_ms(total_underflow_ms),
            "largest_underflow_ms": round_ms(largest_underflow_ms),
        },
        "errors": {
            "total_observed": observed_errors,
            "by_type": dict(sorted(error_counts.items())),
            "empty_audio_delta_count": empty_audio_deltas,
        },
        "analysis_warnings": analysis_warnings,
    }


def format_summary(summary: dict[str, Any]) -> str:
    """Create a concise, copy-pasteable English summary for terminal output."""
    configuration = summary["configuration"]
    completion = summary["completion"]
    latency = summary["latency"]
    frames = summary["frame_delivery"]
    playout = summary["simulated_playout"]
    errors = summary["errors"]
    warnings = summary["analysis_warnings"]
    ttfb = latency["ttfb_ms"]
    rtf = latency["rtf"]
    gaps = frames["raw_interframe_arrival_gap_ms"]
    largest_gap = frames["largest_raw_interframe_arrival_gap"]
    largest_gap_context = (
        "n/a"
        if largest_gap is None
        else (
            f"call {largest_gap['call_id']}, sequences "
            f"{largest_gap['previous_audio_delta_sequence']}->{largest_gap['next_audio_delta_sequence']}"
        )
    )
    return "\n".join(
        [
            "TTS WebSocket Capture Summary",
            f"Load: {configuration['requested_calls']} calls at C={configuration['concurrency']} (prewarm: {'on' if configuration['prewarm'] else 'off'})",
            f"Completion: {completion['audio_done']}/{configuration['requested_calls']} audio.done ({completion['audio_done_rate']:.1%})",
            "TTFB (text.done -> first non-empty audio): "
            f"p50 {display_metric(ttfb['p50'], ' ms')} | p95 {display_metric(ttfb['p95'], ' ms')} | max {display_metric(ttfb['max'], ' ms')}",
            "RTF (last audio arrival / generated audio): "
            f"p50 {display_metric(rtf['p50'])} | p95 {display_metric(rtf['p95'])} | max {display_metric(rtf['max'])}",
            "Raw inter-frame arrival gap (consecutive non-empty audio.delta; excludes TTFB and frame duration): "
            f"p50 {display_metric(gaps['p50'], ' ms')} | p95 {display_metric(gaps['p95'], ' ms')} | max {display_metric(gaps['max'], ' ms')} ({largest_gap_context})",
            "Audio payload duration per frame: "
            f"p50 {display_metric(frames['audio_ms_per_frame']['p50'], ' ms')} | p95 {display_metric(frames['audio_ms_per_frame']['p95'], ' ms')} | max {display_metric(frames['audio_ms_per_frame']['max'], ' ms')}",
            "Simulated underflow "
            f"(TTFB excluded; {playout['initial_buffer_ms']} ms intentional prebuffer): {playout['calls_with_underflow']}/{configuration['requested_calls']} calls | {playout['underflow_events']} events | largest {display_metric(playout['largest_underflow_ms'], ' ms')} | total {display_metric(playout['total_underflow_ms'], ' ms')}",
            f"Errors: {errors['total_observed']} observed | empty audio.delta frames: {errors['empty_audio_delta_count']} | by type: {errors['by_type'] or 'none'}",
            f"Analysis warnings: {len(warnings)}"
            + (f" | {warnings[0]}" if warnings else ""),
        ]
    )


async def run(argv: list[str] | None = None) -> int:
    config = load_config(parser().parse_args(argv))
    run_id = uuid.uuid4().hex
    writer = JsonlWriter(config.output)
    limiter = asyncio.Semaphore(config.concurrency)
    statuses: list[str] = []
    call_ids = {
        number: f"call-{number:04d}-{uuid.uuid4().hex[:8]}"
        for number in range(1, config.calls + 1)
    }
    try:
        if config.prewarm and config.calls != config.concurrency:
            raise ValueError(
                "--prewarm requires --calls to equal --concurrency so the measured load window is consistent."
            )
        await writer.write(
            {
                "event": "run.started",
                "run_id": run_id,
                "started_at_utc": utc_now(),
                "calls": config.calls,
                "concurrency": config.concurrency,
                "prewarm": config.prewarm,
                "event_timeout_ms": None
                if config.event_timeout_s is None
                else round(config.event_timeout_s * 1000),
                "connect_timeout_ms": None
                if config.connect_timeout_s is None
                else round(config.connect_timeout_s * 1000),
                "include_delta_base64": config.include_delta_base64,
                "pcm_format": {
                    "sample_rate": SAMPLE_RATE,
                    "channels": 1,
                    "sample_width_bytes": PCM_BYTES_PER_SAMPLE,
                },
            }
        )
        timeout = aiohttp.ClientTimeout(
            total=None, sock_connect=config.connect_timeout_s
        )
        connector = aiohttp.TCPConnector(limit=config.concurrency)
        async with aiohttp.ClientSession(
            timeout=timeout, connector=connector
        ) as session:
            prewarmed: dict[int, aiohttp.ClientWebSocketResponse | None] = {
                number: None for number in call_ids
            }
            if config.prewarm:
                prewarm_started_at = time.perf_counter()
                results = await asyncio.gather(
                    *[
                        prewarm_connection(
                            session=session,
                            config=config,
                            writer=writer,
                            run_id=run_id,
                            call_number=number,
                            call_id=call_ids[number],
                        )
                        for number in call_ids
                    ],
                    return_exceptions=True,
                )
                failures = [item for item in results if isinstance(item, BaseException)]
                if failures:
                    for item in results:
                        if (
                            isinstance(item, aiohttp.ClientWebSocketResponse)
                            and not item.closed
                        ):
                            await item.close()
                    raise RuntimeError(
                        f"Prewarm failed for {len(failures)}/{config.calls} connections."
                    )
                prewarmed = {number: results[number - 1] for number in call_ids}
                await writer.write(
                    {
                        "event": "run.prewarm_finished",
                        "run_id": run_id,
                        "finished_at_utc": utc_now(),
                        "prewarm_duration_ms": elapsed_ms(prewarm_started_at),
                        "connections_ready": len(prewarmed),
                    }
                )
            statuses = await asyncio.gather(
                *[
                    capture_call(
                        session=session,
                        config=config,
                        writer=writer,
                        run_id=run_id,
                        call_number=call_number,
                        call_id=call_ids[call_number],
                        limiter=limiter,
                        prewarmed_ws=prewarmed[call_number],
                    )
                    for call_number in call_ids
                ]
            )
        await writer.write(
            {
                "event": "run.finished",
                "run_id": run_id,
                "finished_at_utc": utc_now(),
                "statuses": {
                    status: statuses.count(status) for status in sorted(set(statuses))
                },
            }
        )
    except BaseException as exc:
        await writer.write(
            {
                "event": "run.failed",
                "run_id": run_id,
                "failed_at_utc": utc_now(),
                "error_type": type(exc).__name__,
                "error": str(exc),
            }
        )
        raise
    finally:
        writer.close()
        try:
            summary = summarize_capture(config, run_id)
            config.summary_output.parent.mkdir(parents=True, exist_ok=True)
            config.summary_output.write_text(
                json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
            print(format_summary(summary))
            print(f"Summary JSON: {config.summary_output}")
        except Exception as exc:
            print(
                f"Unable to generate summary JSON: {type(exc).__name__}: {exc}",
                file=sys.stderr,
            )
    return 0 if statuses and all(status == "audio_done" for status in statuses) else 1


def main() -> None:
    raise SystemExit(asyncio.run(run()))


if __name__ == "__main__":
    main()
