"""Zonos-gateway TTS client + PAD→emotion-dial mapping.

Slice 2 of docs/contracts/donut_voiced_interview.contract.md. This module is the
SINGLE swap seam for voice synthesis: the `/api/tts` route in web/server.py is
its only caller. Direct coupling to the Zonos gateway (DEC-1) buys the emotion
dials that the swappable `ext-tts` LiteLLM alias drops — the whole point is
affect-driven voice (DEC-7). If we ever move off Zonos, this is the swap point.

Foot-guns (verified live 2026-08-02):
  - DEC-3: response_format is ALWAYS "wav". `mp3`/`opus` are accepted but the
    gateway silently returns mislabeled PCM (no encoder wired) — never request them.
  - Use the gateway :8890, NOT the engine :1920 (rep-penalty bug pads silence).
"""

from __future__ import annotations

import re
import sys
from collections.abc import AsyncIterator, Mapping
from dataclasses import dataclass

import httpx

# DEC-1: the Zonos gateway (irv-ml1 :8890) — NOT the engine :1920, NOT the
# swappable ext-tts alias (which drops the emotion dials). Overridable per
# deployment via app.state.tts_url (RATATOSKR_TTS_URL) — the swap seam + tests.
ZONOS_TTS_URL = "http://10.100.79.3:8890/v1/audio/speech"

# DEC-7: fixed emotion strength when PAD-driven (tunable; the gateway scales the
# valence/arousal push by this).
_DEFAULT_EMOTION_STRENGTH = 1.0

# DEC-9: pin English conditioning. Zonos is multilingual; without an explicit
# `language` it drifts into other-language phonemes / gibberish on names, numerics,
# and long spans (operator report 2026-08-02). The gateway accepts an ISO-ish
# `language` code (verified live against /v1/audio/speech).
_TTS_LANGUAGE = "en-us"


class TtsUnavailable(Exception):
    """The Zonos gateway failed, was unreachable, or returned a non-wav body.

    The caller degrades (INV-TTS-4): logs + skips audio; the turn/transcript is
    never blocked or failed on a synthesis error.
    """

    def __init__(self, message: str = "", *, status: int = 0) -> None:
        super().__init__(message)
        self.status = status
        self.message = message


@dataclass(frozen=True)
class PadState:
    """Live PAD read off the affect_update SSE `current` snapshot (the console
    already consumes it, DEC-7). `dominance` is carried for completeness but
    unused by the emotion dials — Zonos exposes valence + arousal only."""

    pleasure: float
    arousal: float
    dominance: float = 0.0

    @classmethod
    def from_obj(cls, obj: object) -> PadState | None:
        """Open-world parse of the browser-sent `pad`. Returns None on a missing
        or malformed value (pad_to_dials then degrades to a neutral read). Never
        raises — the wire is untrusted (INV-TTS-4 / degrade-not-crash)."""
        if not isinstance(obj, Mapping):
            return None
        try:
            return cls(
                pleasure=float(obj["pleasure"]),
                arousal=float(obj["arousal"]),
                dominance=float(obj.get("dominance", 0.0)),
            )
        except (KeyError, TypeError, ValueError, ArithmeticError):
            # ArithmeticError covers OverflowError — float() of a huge JSON integer
            # literal (e.g. a 400-digit number) raises it; the wire is untrusted, so a
            # malformed pad degrades to a neutral read rather than 500ing /api/tts.
            return None


@dataclass(frozen=True)
class EmotionDials:
    """Zonos emotion dials (DEC-7). Fields mirror the gateway's /v1/dials surface
    (verified live 2026-08-02). `emotion_enabled=False` is a neutral read: the
    other fields are omitted from the POST body, so the gateway uses its default
    voice emotion."""

    emotion_enabled: bool = False
    emotion_valence: float = 0.0
    emotion_arousal: float = 0.0
    emotion_strength: float = _DEFAULT_EMOTION_STRENGTH

    def to_body(self) -> dict:
        """The dial fields for the gateway POST body. Emitted ONLY when enabled;
        a neutral read contributes nothing (the gateway falls to its default)."""
        if not self.emotion_enabled:
            return {}
        return {
            "emotion_enabled": True,
            "emotion_valence": self.emotion_valence,
            "emotion_arousal": self.emotion_arousal,
            "emotion_strength": self.emotion_strength,
        }


def _clamp(x: float, lo: float, hi: float) -> float:
    """Clamp to [lo, hi], NaN-safe: a NaN axis (a dead/malformed signal) → 0.0
    rather than propagating through max/min into the gateway."""
    if x != x:  # NaN
        return 0.0
    return max(lo, min(hi, x))


def pad_to_dials(pad: PadState | None) -> EmotionDials:
    """Map live PAD → Zonos emotion dials (DEC-7 / FN pad_to_dials).

    Total: any input (None / finite / out-of-range / NaN / inf / a non-PadState
    object) → valid dials, never raises. None/absent/malformed PAD → a neutral read
    (emotion_enabled=False)."""
    if not isinstance(pad, PadState):
        # None, or any non-PadState the declared surface doesn't cover — neutral read.
        return EmotionDials(emotion_enabled=False)
    return EmotionDials(
        emotion_enabled=True,
        emotion_valence=_clamp(pad.pleasure, -1.0, 1.0),
        emotion_arousal=_clamp(pad.arousal, -1.0, 1.0),
        emotion_strength=_DEFAULT_EMOTION_STRENGTH,
    )


def gateway_body(text: str, voice: str, dials: EmotionDials) -> dict:
    """The Zonos gateway POST body — response_format is ALWAYS "wav" (DEC-3);
    `language` is pinned to en-us (DEC-9) so the multilingual model stays in English."""
    return {
        "input": text,
        "voice": voice,
        "response_format": "wav",
        "language": _TTS_LANGUAGE,
        **dials.to_body(),
    }


# DEC-10: long-form chunking. The Zonos model hard-caps ONE synthesis at max_tokens=6144
# = 71.2s of audio (6144 / 86.3 Hz codec frame rate; infra-ops verified 2026-08-02, and the
# gateway 400s above 6144 — the model's architectural sequence limit, unraisable). Anything
# longer must be chunk-and-concatenate client-side. We target 75% of the cap per chunk
# (operator call: greedy for prosody — fewer, fuller chunks, fewer seams), leaving the other
# 25% as headroom for char->audio-seconds variance: the cap is on TOKENS (6144 = 71.2s ALWAYS,
# a fixed codec-frame constant independent of delivery — infra-ops), but our budget is in CHARS
# (a proxy), and slower/expressive delivery stretches a given char-count into more audio-seconds
# (more tokens). ~14 chars/s at neutral rate => ~747-char budget (~53s).
_TTS_MODEL_CAP_SECONDS = 71.2
_TTS_CHUNK_TARGET_FRACTION = 0.75
_TTS_CHARS_PER_SEC = 14.0
_TTS_CHUNK_CHAR_BUDGET = int(
    _TTS_MODEL_CAP_SECONDS * _TTS_CHUNK_TARGET_FRACTION * _TTS_CHARS_PER_SEC
)

_PARA_SPLIT = re.compile(r"\n\s*\n+")  # blank-line paragraph boundary
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?])\s+")  # infra-ops recipe: split after . ! ?
_CLAUSE_BOUNDARIES = (", ", "; ", ": ")


def _hard_wrap(s: str, budget: int) -> list[str]:
    """Sub-split a single over-budget sentence into <= budget pieces, preferring a clause
    boundary (`, ; :`), then any space; hard-cuts mid-word ONLY as a last resort (input
    with no usable boundary at all). Never returns a piece longer than `budget`."""
    out: list[str] = []
    while len(s) > budget:
        window = s[:budget]
        cut = max((window.rfind(c) for c in _CLAUSE_BOUNDARIES), default=-1)
        if cut > 0:
            cut += 1  # keep the delimiter char with the head; break just after it
        else:
            cut = window.rfind(" ")
        if cut <= 0:
            cut = budget  # no boundary in range — last-resort hard cut
        out.append(s[:cut].strip())
        s = s[cut:].strip()
    if s:
        out.append(s)
    return out


def _greedy_pack(units: list[str], budget: int, join: str) -> list[str]:
    """Greedily pack pre-sized units (each <= budget) into <= budget chunks joined by
    `join` — as few and as full as possible, so prosody flows across the fewest seams."""
    chunks: list[str] = []
    cur = ""
    for u in units:
        if cur and len(cur) + len(join) + len(u) > budget:
            chunks.append(cur)
            cur = u
        else:
            cur = f"{cur}{join}{u}" if cur else u
    if cur:
        chunks.append(cur)
    return chunks


def chunk_text(text: str, budget: int = _TTS_CHUNK_CHAR_BUDGET) -> list[str]:
    """Split `text` into synthesis chunks each <= `budget` chars (DEC-10 / FN chunk_text).

    Paragraph-first (operator call): whole paragraphs greedily pack together while they fit,
    so seams land on blank-line boundaries where a natural pause already belongs. A paragraph
    over budget flushes the pending whole-paragraph run, then falls back to sentence packing;
    a sentence over budget falls back to clause/word sub-splitting (never mid-word unless the
    input has no boundary at all). Every returned chunk is non-empty and <= budget;
    empty/whitespace input -> []."""
    # A non-positive budget would make _hard_wrap spin forever (cut=budget<=0 => no
    # forward progress). Production always passes the 747 default; clamp a misuse to a
    # sane floor rather than hang the event loop (degrade-not-crash, module-wide ethos).
    budget = max(1, budget)
    text = (text or "").strip()
    if not text:
        return []
    chunks: list[str] = []
    pending: list[str] = []  # whole paragraphs (each <= budget) awaiting a greedy pack

    def flush_pending() -> None:
        if pending:
            chunks.extend(_greedy_pack(pending, budget, join="\n\n"))
            pending.clear()

    for para in _PARA_SPLIT.split(text):
        para = para.strip()
        if not para:
            continue
        if len(para) <= budget:
            pending.append(para)
            continue
        # Oversized paragraph: emit the accumulated whole-paragraph chunks first (don't
        # merge a mid-paragraph fragment across the blank-line boundary), then sentence-pack.
        flush_pending()
        sentences: list[str] = []
        for sent in _SENTENCE_SPLIT.split(para):
            sent = sent.strip()
            if not sent:
                continue
            if len(sent) <= budget:
                sentences.append(sent)
            else:
                sentences.extend(_hard_wrap(sent, budget))
        chunks.extend(_greedy_pack(sentences, budget, join=" "))
    flush_pending()
    return chunks


async def tts_stream(
    text: str,
    *,
    voice: str,
    dials: EmotionDials,
    client: httpx.AsyncClient,
    url: str = ZONOS_TTS_URL,
) -> AsyncIterator[bytes]:
    """Open the gateway's CHUNKED stream and yield WAV bytes as they synthesize.

    The gateway emits a streaming int16 WAV (RIFF/data sizes = 0xFFFFFFFF placeholders)
    over `transfer-encoding: chunked`, TTFB ~0.44s vs ~7s total (infra-ops verified) —
    designed to be played progressively by a browser <audio>. So we PROXY THE CHUNKS
    STRAIGHT THROUGH: never buffer, never rewrite the header (that would force us to wait
    for the whole clip and defeat the streaming — the bug this replaces).

    Raises TtsUnavailable if the gateway open is a non-200 or the connect/transport
    fails — BEFORE the first chunk, so the endpoint can still return a 503 (INV-TTS-4).
    A drop mid-stream just ends the generator (the browser has already played the head).
    """
    assert text, "tts_stream: text must be non-empty (the endpoint guards this)"
    try:
        async with client.stream("POST", url, json=gateway_body(text, voice, dials)) as resp:
            if resp.status_code != 200:
                raise TtsUnavailable(
                    f"gateway status {resp.status_code}", status=resp.status_code
                )
            async for chunk in resp.aiter_bytes():
                yield chunk
    except httpx.RequestError as exc:
        raise TtsUnavailable(f"gateway transport failure: {exc}") from exc


# The streaming WAV header is 44 bytes (RIFF 12 + fmt 24 + data 8; infra-ops 2026-08-02).
# A body that has not produced a `data` chunk id within this window is not the WAV we expect,
# so we stop scanning rather than buffer unbounded / match a `data` byte-run deep in PCM.
_WAV_HEADER_SCAN_LIMIT = 1024


async def _pcm_after_header(stream: AsyncIterator[bytes]) -> AsyncIterator[bytes]:
    """Yield only the PCM payload of a streaming WAV — everything AFTER the `data` chunk
    id+size (8 bytes). Used for chunks 2..N so their RIFF header isn't buried mid-stream
    (infra-ops 2026-08-02: one header only, or the concatenated file corrupts). The `data`
    marker can straddle two network reads, so accumulate until it's found + 8 bytes land.

    Raises TtsUnavailable if the body has no `data` chunk within the header window, or ends
    before one completes — a malformed / non-WAV 200 body. tts_stream_long catches that for a
    chunks-2..N failure and degrades (INV-TTS-4) rather than emitting header bytes as PCM or
    buffering the whole body forever."""
    acc = bytearray()
    header_done = False
    async for b in stream:
        if header_done:
            yield b
            continue
        acc.extend(b)
        di = acc.find(b"data")
        if di >= 0 and di + 8 <= len(acc):
            header_done = True
            tail = bytes(acc[di + 8:])
            acc = bytearray()
            if tail:
                yield tail
        elif len(acc) > _WAV_HEADER_SCAN_LIMIT:
            # No `data` chunk within a sane header window — not the streaming WAV we expect.
            raise TtsUnavailable("chunk 2+ body has no WAV data chunk within the header window")
    if not header_done:
        # Stream ended before a complete `data` header — truncated / empty chunk body.
        raise TtsUnavailable("chunk 2+ body ended before the WAV data chunk")


async def tts_stream_long(
    text: str,
    *,
    voice: str,
    dials: EmotionDials,
    client: httpx.AsyncClient,
    url: str = ZONOS_TTS_URL,
    budget: int = _TTS_CHUNK_CHAR_BUDGET,
) -> AsyncIterator[bytes]:
    """Synthesize arbitrarily long `text` as ONE continuous int16-PCM stream by chunking it
    under the model's 71.2s cap (DEC-10) and concatenating. Chunk 1 streams verbatim (its
    WAV header + PCM); chunks 2..N stream PCM-only (header stripped, `_pcm_after_header`) so
    the browser decodes one gapless stream after a single leading header. Identical
    voice+dials on every chunk so delivery stays uniform across the seams.

    A failure BEFORE the first byte (an OPEN failure) propagates as TtsUnavailable so the
    endpoint peek turns it into a 503 (INV-TTS-4) — nothing is committed yet. A failure AFTER
    bytes have already streamed — a MID-STREAM drop, on chunk 0 or a later chunk, after the
    200 is committed — degrades: drop the tail, keep what played, log, and NEVER raise into
    the committed StreamingResponse (which would corrupt it with an ASGI trace). The pivot is
    `yielded_any`, not the chunk index — a chunk-0 mid-stream drop is a committed-200 failure
    too, not an open failure."""
    chunks = chunk_text(text, budget)
    yielded_any = False
    for i, chunk in enumerate(chunks):
        stream = tts_stream(chunk, voice=voice, dials=dials, client=client, url=url)
        src = stream if i == 0 else _pcm_after_header(stream)
        try:
            async for b in src:
                yielded_any = True
                yield b
        except TtsUnavailable as exc:
            if not yielded_any:
                raise  # open failure, pre-commit → endpoint peek → 503
            # committed-200 mid-stream failure → degrade + log (INV-TTS-4 "logs + skips").
            sys.stderr.write(
                f'{{"kind":"tts_degrade","event":"chunk_failed","chunk_index":{i},'
                f'"chunks_total":{len(chunks)},"exc":"{type(exc).__name__}"}}\n'
            )
            return
