"""chatterbox-fast — streaming TTS server (Phase 1 MVP).

A lean FastAPI server on the ChatterboxTurboTTS library that streams audio using
the adaptive buffer-ratchet scheduler in ``scheduler.py`` (the meat). Sub-second
time-to-first-audio while keeping turbo's full quality; workload is single-stream
interactive (see docs/design/chatterbox-fast-plan.md).

Endpoints:
  POST /tts     — StreamingResponse of audio chunks (raw PCM s16le default).
  GET  /health  — model/voice readiness.

Config via env (all optional; sane dev defaults):
  CBF_MODEL_DEVICE   cuda | cuda:0 | cpu                  (default: cuda)
  CBF_VOICES_DIR     dir of predefined voice wavs          (default: /refs)
  CBF_DEFAULT_VOICE  default reference wav path/name       (default: first wav in dir)
  CBF_BIND / CBF_PORT  uvicorn bind                        (default: 0.0.0.0:8197)
"""

from __future__ import annotations

import logging
import os
import struct
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Iterator, Literal

import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from chatterbox_fast.scheduler import ChunkConfig, ChunkResult, stream_chunks

log = logging.getLogger("chatterbox-fast")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# ── config ────────────────────────────────────────────────────────────────

DEVICE = os.environ.get("CBF_MODEL_DEVICE", "cuda")
VOICES_DIR = Path(os.environ.get("CBF_VOICES_DIR", "/refs"))
DEFAULT_VOICE_ENV = os.environ.get("CBF_DEFAULT_VOICE")
BIND = os.environ.get("CBF_BIND", "0.0.0.0")
PORT = int(os.environ.get("CBF_PORT", "8197"))

# Perf levers (plan §4 Phase 2). TF32 + flash/mem-efficient SDPA are low-risk on
# Ampere and free — default ON. Measured 2026-06-02: they do NOT move TTFA, which
# is bound by the sequential AR token decode (T3 Llama at batch-1), not matmul
# throughput. bf16 (the lever that WOULD help batch-1 decode) is DEFERRED: turbo
# loads fp32 and from_pretrained() exposes no dtype arg, so bf16 needs whole-model
# casting incl. the speaker-conditioning path and the dtype-sensitive vocoder —
# real surgery + quality risk for a TTFA gain not currently needed (~0.5s is fine).
PERF_TF32 = os.environ.get("CBF_TF32", "1") == "1"
PERF_SDPA_FLASH = os.environ.get("CBF_SDPA_FLASH", "1") == "1"

# Turbo sampling knobs validated in the spike (plan §2). CFG / exaggeration /
# min_p are ignored by turbo (it warns, harmless).
WARMUP_TEXT = "Warming up the streaming engine."


def _predefined_wavs() -> list[Path]:
    """Predefined voice wavs in VOICES_DIR, excluding `_`-prefixed scratch files
    (bench/A-B outputs land as `_*.wav` in the same dir)."""
    if not VOICES_DIR.is_dir():
        return []
    return sorted(p for p in VOICES_DIR.glob("*.wav") if not p.name.startswith("_"))


def _setup_perf() -> None:
    """Apply the safe, low-risk speed levers before model load."""
    if PERF_TF32:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        log.info("perf: TF32 matmul/cudnn enabled")
    if PERF_SDPA_FLASH and DEVICE.startswith("cuda"):
        try:
            torch.backends.cuda.enable_flash_sdp(True)
            torch.backends.cuda.enable_mem_efficient_sdp(True)
            log.info("perf: flash + mem-efficient SDPA enabled")
        except Exception as e:  # pragma: no cover - backend-dependent
            log.warning("perf: SDPA toggle failed: %s", e)


# ── model holder ──────────────────────────────────────────────────────────


class Engine:
    """Owns the single ChatterboxTurboTTS instance + a generation lock.

    The workload is single-stream interactive, but prepare_conditionals mutates
    model state, so all generation is serialized under one lock. Concurrent
    multi-voice serving is a Phase-2 concern.
    """

    def __init__(self) -> None:
        self.model = None
        self.sr = 24000
        self.lock = threading.Lock()
        self._current_voice: str | None = None
        self.default_voice: str | None = None

    def load(self) -> None:
        from chatterbox.tts_turbo import ChatterboxTurboTTS

        _setup_perf()
        log.info("loading ChatterboxTurboTTS on %s …", DEVICE)
        t0 = time.perf_counter()
        self.model = ChatterboxTurboTTS.from_pretrained(device=DEVICE)
        self.sr = int(getattr(self.model, "sr", 24000))
        self._log_model_dtype()
        self.default_voice = self._discover_default_voice()
        if self.default_voice:
            self._prepare(self.default_voice)
            self._warm()
        log.info(
            "model ready in %.1fs (sr=%d, default_voice=%s)",
            time.perf_counter() - t0, self.sr, self.default_voice,
        )

    def _log_model_dtype(self) -> None:
        for name in ("t3", "s3gen", "model"):
            sub = getattr(self.model, name, None)
            try:
                if sub is not None:
                    dt = next(sub.parameters()).dtype
                    log.info("dtype[%s]=%s", name, dt)
            except (StopIteration, AttributeError):
                pass

    def _discover_default_voice(self) -> str | None:
        if DEFAULT_VOICE_ENV:
            # Resolve a bare name ("glados_25s") the same way request-time voices
            # are resolved — append .wav and look under VOICES_DIR.
            return self.resolve_voice(DEFAULT_VOICE_ENV)
        wavs = _predefined_wavs()
        return str(wavs[0]) if wavs else None

    def resolve_voice(self, voice: str | None) -> str:
        if not voice:
            if not self.default_voice:
                raise HTTPException(503, "no default voice; set CBF_DEFAULT_VOICE")
            return self.default_voice
        p = Path(voice)
        if p.is_absolute() and p.exists():
            return str(p)
        cand = VOICES_DIR / (voice if voice.endswith(".wav") else f"{voice}.wav")
        if cand.exists():
            return str(cand)
        raise HTTPException(404, f"unknown voice {voice!r}")

    def _prepare(self, voice_path: str, exaggeration: float = 0.5) -> None:
        if voice_path == self._current_voice:
            return
        log.info("prepare_conditionals(%s)", voice_path)
        self.model.prepare_conditionals(voice_path, exaggeration=exaggeration, norm_loudness=True)
        self._current_voice = voice_path

    def _warm(self) -> None:
        with torch.inference_mode():
            self.model.generate(WARMUP_TEXT, repetition_penalty=1.2, top_p=0.95,
                                 temperature=0.8, top_k=1000)
        if DEVICE.startswith("cuda"):
            torch.cuda.synchronize()

    def generate(self, text: str, knobs: "TTSRequest") -> tuple[torch.Tensor, float]:
        """Synthesize ``text`` → (wav tensor [1,T], audio_seconds). CUDA-synced
        so the caller's clock delta is honest gen time."""
        with torch.inference_mode():
            wav = self.model.generate(
                text,
                repetition_penalty=knobs.repetition_penalty,
                top_p=knobs.top_p,
                temperature=knobs.temperature,
                top_k=knobs.top_k,
            )
        if DEVICE.startswith("cuda"):
            torch.cuda.synchronize()
        audio_sec = wav.shape[-1] / self.sr
        return wav, audio_sec


engine = Engine()


@asynccontextmanager
async def lifespan(app: FastAPI):
    engine.load()
    yield


app = FastAPI(title="chatterbox-fast", lifespan=lifespan)


# ── request / audio encoding ──────────────────────────────────────────────


class TTSRequest(BaseModel):
    text: str
    voice: str | None = None
    format: Literal["pcm", "wav"] = "pcm"
    stream: bool = True  # False ⇒ whole-text one-shot (for A/B vs streaming)
    exaggeration: float = 0.5
    temperature: float = 0.8
    top_p: float = 0.95
    top_k: int = 1000
    repetition_penalty: float = 1.2
    seed: int = 0  # 0 ⇒ random; a fixed seed repeats a one-shot take (see note below)
    # Scheduler overrides (None ⇒ ChunkConfig defaults).
    margin: float | None = Field(default=None)
    margin_first: float | None = Field(default=None)
    rtf_prior: float | None = Field(default=None)


def _pcm16(wav: torch.Tensor) -> bytes:
    a = wav.detach().to(torch.float32).clamp_(-1.0, 1.0).cpu().numpy().reshape(-1)
    return (a * 32767.0).astype("<i2").tobytes()


def _wav_header(sr: int, data_len: int | None = None) -> bytes:
    """WAV header. data_len=None → streaming (length unknown, 0xFFFFFFFF sizes,
    player reads to EOF); an int → correct RIFF/data sizes for a complete file."""
    data_size = 0xFFFFFFFF if data_len is None else data_len
    riff_size = 0xFFFFFFFF if data_len is None else 36 + data_len
    return (
        b"RIFF" + struct.pack("<I", riff_size) + b"WAVE"
        + b"fmt " + struct.pack("<IHHIIHH", 16, 1, 1, sr, sr * 2, 2, 16)
        + b"data" + struct.pack("<I", data_size)
    )


def _chunk_config(req: TTSRequest) -> ChunkConfig:
    cfg = ChunkConfig()
    if req.margin is not None:
        cfg.margin = req.margin
    if req.margin_first is not None:
        cfg.margin_first = req.margin_first
    if req.rtf_prior is not None:
        cfg.rtf_prior = req.rtf_prior
    return cfg


# ── endpoints ─────────────────────────────────────────────────────────────


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok" if engine.model is not None else "loading",
        "sr": engine.sr,
        "device": DEVICE,
        "default_voice": engine.default_voice,
        "voices_dir": str(VOICES_DIR),
    }


@app.get("/voices")
def voices() -> dict:
    """Predefined voices = the *.wav stems in CBF_VOICES_DIR. Clone refs are
    passed per-request as an absolute path, so they're not listed here."""
    names = [p.stem for p in _predefined_wavs()]
    default = Path(engine.default_voice).stem if engine.default_voice else None
    return {"voices": names, "default": default}


@app.post("/tts")
def tts(req: TTSRequest) -> StreamingResponse:
    if engine.model is None:
        raise HTTPException(503, "model still loading")
    if not req.text.strip():
        raise HTTPException(400, "empty text")

    voice_path = engine.resolve_voice(req.voice)
    cfg = _chunk_config(req)
    media = "audio/wav" if req.format == "wav" else "application/octet-stream"

    def body() -> Iterator[bytes]:
        # One request holds the lock for its whole stream (single-stream
        # workload); concurrent callers queue rather than corrupt conditionals.
        with engine.lock:
            engine._prepare(voice_path, exaggeration=req.exaggeration)
            # Seed once per request (under the lock). One-shot is then reproducible
            # for a fixed seed + params; streaming is NOT — adaptive-chunk boundaries
            # depend on live-measured RTF (wall-clock), so chunk splits vary run to run.
            if req.seed:
                torch.manual_seed(req.seed)
                if DEVICE.startswith("cuda"):
                    torch.cuda.manual_seed_all(req.seed)
            t_req = time.perf_counter()
            first_audio_ms: float | None = None

            if not req.stream:
                # One-shot: the full length is known, so emit a correct-sized WAV
                # header (a buffered consumer wants well-formed sizes).
                wav, audio_sec = engine.generate(req.text, req)
                first_audio_ms = (time.perf_counter() - t_req) * 1000
                log.info("oneshot: %.0fms gen, %.2fs audio", first_audio_ms, audio_sec)
                pcm = _pcm16(wav)
                if req.format == "wav":
                    yield _wav_header(engine.sr, len(pcm))
                yield pcm
                return

            # Streaming: length is unknown up front → open-ended WAV header.
            if req.format == "wav":
                yield _wav_header(engine.sr)

            def _gen(text: str) -> tuple[torch.Tensor, float]:
                return engine.generate(text, req)

            total_audio = 0.0
            for r in stream_chunks(req.text, generate=_gen, clock=time.perf_counter, cfg=cfg):
                if first_audio_ms is None:
                    first_audio_ms = (time.perf_counter() - t_req) * 1000
                total_audio += r.audio_sec
                _log_chunk(r, first_audio_ms)
                yield _pcm16(r.audio)
            log.info("stream done: ttfa=%.0fms total_audio=%.2fs", first_audio_ms or 0, total_audio)

    return StreamingResponse(body(), media_type=media)


def _log_chunk(r: ChunkResult, ttfa_ms: float) -> None:
    tag = " STARVED" if r.starved else ""
    if r.index == 0:
        log.info("chunk 0: ttfa=%.0fms gen=%.0fms audio=%.2fs rtf=%.2f%s",
                 ttfa_ms, r.gen_time * 1000, r.audio_sec, r.rtf, tag)
    else:
        log.info("chunk %d: gen=%.0fms audio=%.2fs buf=%.2f→%.2f rtf=%.2f%s",
                 r.index, r.gen_time * 1000, r.audio_sec,
                 r.buffer_before, r.buffer_after, r.rtf, tag)


def main() -> None:
    """Console entrypoint (`chatterbox-fast`) and Docker CMD."""
    import uvicorn

    uvicorn.run(app, host=BIND, port=PORT)


if __name__ == "__main__":
    main()