#!/usr/bin/env python3
"""OpenAI-compatible TTS wrapper around chatterbox-tts, for Open Notebook.

Implements exactly the endpoint esperanto's OpenAICompatibleTextToSpeechModel calls:
    POST /audio/speech  {model, voice, input, response_format} -> raw audio bytes

Voice cloning: `voice` selects a named reference clip from VOICES_DIR (each
<name>.wav becomes voice "<name>"). Unknown/omitted voices fall back to
"default". A recognized language code (see chatterbox_cli_v4.SUPPORTED_LANGS)
is still accepted as `voice` for the old behaviour: no cloning, built-in voice
for that language. All cloned voices are generated in German (CLONE_LANG) —
this deployment only has German reference clips; extend VOICE_LANG_OVERRIDES
below if you add clips in other languages.

Start:
    ~/miniforge3/envs/chatterbox/bin/python tts_server.py

Env vars:
    TTS_HOST (default 0.0.0.0), TTS_PORT (default 8901)
    CUDA_VISIBLE_DEVICES should be set by the caller (e.g. "2") to keep GPU 0/1 free.
"""
from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import threading
from pathlib import Path

sys.path.insert(0, str(Path.home() / "chatterbox-tts-cli"))
import chatterbox_cli_v4 as tts  # noqa: E402

import torch
import torchaudio as ta
from fastapi import FastAPI, HTTPException, Response
from pydantic import BaseModel

app = FastAPI(title="Chatterbox TTS (OpenAI-compatible, voice cloning)", version="2.0")

_DEVICE = tts.get_device(None)
_model_cache: dict[str, tuple] = {}

# Serializes model load + inference. FastAPI runs sync handlers in a threadpool, so
# without this every concurrent request generates in parallel on the same GPU and the
# activation memory multiplies — podcast_creator sends TTS_BATCH_SIZE (default 5) clips
# at once, which OOM'd this server (peak 16.7 GB) since GPU 2 is shared with the STT
# server and the chatterbox MCP service. Generation is GPU-bound, so serializing costs
# no real throughput; it just bounds peak VRAM to a single clip.
_GPU_LOCK = threading.Lock()

VOICES_DIR = Path(__file__).parent / "voices"
CLONE_LANG = "de"  # language used for every cloned reference voice
VOICE_LANG_OVERRIDES: dict[str, str] = {}  # e.g. {"john": "en"} if you add an English clip


def _discover_voices() -> dict[str, Path]:
    if not VOICES_DIR.is_dir():
        return {}
    return {p.stem: p for p in sorted(VOICES_DIR.glob("*.wav"))}


VOICES = _discover_voices()


def _get_model(lang: str):
    key = "en" if lang == "en" else "multi"
    if key not in _model_cache:
        _model_cache[key] = tts.load_model(lang, _DEVICE, t3_model="v3")
    return _model_cache[key]


def _resolve_voice(voice: str) -> tuple[str, str | None]:
    """Map the incoming `voice` string to (language, audio_prompt_path|None)."""
    if voice in VOICES:
        return VOICE_LANG_OVERRIDES.get(voice, CLONE_LANG), str(VOICES[voice])
    if voice in tts.SUPPORTED_LANGS:
        return voice, None
    if "default" in VOICES:
        return CLONE_LANG, str(VOICES["default"])
    return "de", None


class SpeechRequest(BaseModel):
    model: str | None = None
    voice: str = "default"
    input: str
    response_format: str = "mp3"


_FORMAT_CONTENT_TYPE = {
    "mp3": "audio/mpeg",
    "wav": "audio/wav",
    "opus": "audio/opus",
    "flac": "audio/flac",
    "aac": "audio/aac",
}


@app.get("/health")
def health():
    return {"status": "ok", "device": _DEVICE, "voices": sorted(VOICES.keys())}


@app.get("/audio/voices")
def voices():
    cloned = [
        {
            "id": name,
            "name": name,
            "gender": "NEUTRAL",
            "language_code": VOICE_LANG_OVERRIDES.get(name, CLONE_LANG),
            "description": f"Cloned voice from {path.name}",
        }
        for name, path in VOICES.items()
    ]
    langs = [
        {"id": lang, "name": lang, "gender": "NEUTRAL", "language_code": lang,
         "description": "Built-in voice (no cloning)"}
        for lang in sorted(tts.SUPPORTED_LANGS)
    ]
    return {"voices": cloned + langs}


@app.post("/audio/speech")
def speech(req: SpeechRequest):
    lang, voice_path = _resolve_voice(req.voice)

    raw = tts.clean_raw_text(req.input)
    raw_chunks = tts.split_into_sentences(raw, max_len=400)
    chunks = [tts.preprocess_tts_text(c, lang=lang, pronunciation_dict=None) for c in raw_chunks]
    chunks = [c for c in chunks if c.strip()]
    if not chunks:
        raise HTTPException(status_code=422, detail="Kein synthetisierbarer Text übrig.")

    with _GPU_LOCK:
        model, model_kind, sr = _get_model(lang)
        try:
            wavs = []
            for chunk in chunks:
                wavs.append(tts.generate_chunk(model, model_kind, chunk, lang, voice_path))
            final = wavs[0] if len(wavs) == 1 else torch.cat(wavs, dim=-1)
        finally:
            # GPU 2 is shared; hand cached blocks back so a neighbour process can't be
            # starved by fragmentation we're holding on to.
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as wav_tmp:
        wav_path = wav_tmp.name
    ta.save(wav_path, final, sr)

    fmt = req.response_format if req.response_format in _FORMAT_CONTENT_TYPE else "wav"
    try:
        if fmt == "wav":
            audio_bytes = Path(wav_path).read_bytes()
        else:
            out_path = wav_path.replace(".wav", f".{fmt}")
            subprocess.run(
                ["ffmpeg", "-y", "-i", wav_path, out_path],
                check=True, capture_output=True,
            )
            audio_bytes = Path(out_path).read_bytes()
            os.unlink(out_path)
    finally:
        os.unlink(wav_path)

    return Response(content=audio_bytes, media_type=_FORMAT_CONTENT_TYPE[fmt])


if __name__ == "__main__":
    import uvicorn
    print(f"[chatterbox-tts] Registrierte Stimmen: {sorted(VOICES.keys()) or '(keine, nur Sprachcodes)'}")
    uvicorn.run(app, host=os.environ.get("TTS_HOST", "0.0.0.0"),
                port=int(os.environ.get("TTS_PORT", "8901")))