"""IndexTTS-2 — minimal FastAPI wrapper.

Upstream (https://github.com/index-tts/index-tts) ships only a Gradio
webui; the only existing FastAPI fork (csllpr/index-tts-fastapi) targets
v1 and is dormant. We own this wrapper end-to-end.

API surface:
  POST /v1/audio/speech            OpenAI-compat-ish, see SpeechRequest
  GET  /v1/voices                  list speakers + emotion references
  GET  /healthz                    liveness for compose healthcheck

Voice library is a flat directory of WAVs (one file per voice). Emotion
references live in a parallel directory. Both are bind-mounted from the
host so cloned voices survive container recreates.

Emotion control is opt-in and mutually exclusive (audio > vector > text):
  * emotion_voice  — name of a WAV in the emotions dir, used as a SECOND
                     reference whose timbre is ignored but emotion is
                     transferred onto the speaker.
  * emotion_vector — 8-float [happy, angry, sad, afraid, disgusted,
                     melancholic, surprised, calm].
  * emotion_text   — free text; bundled QwenEmotion model derives the
                     vector ("she said excitedly" → joy spike).

Without any of these the speaker WAV's natural emotion is reused.
"""
from __future__ import annotations

import io
import logging
import os
from pathlib import Path
from typing import Iterable, List, Optional

# IndexTTS pins HF_HUB_CACHE at import time (./checkpoints/hf_cache by
# default, see infer_v2.py:4). Override BEFORE the indextts import or
# downloads land in the wrong tree.
os.environ.setdefault(
    "HF_HUB_CACHE",
    os.environ.get("INDEX_TTS_HF_CACHE", "/app/checkpoints/hf_cache"),
)

import numpy as np  # noqa: E402
import soundfile as sf  # noqa: E402
from fastapi import FastAPI, HTTPException  # noqa: E402
from fastapi.responses import Response, StreamingResponse  # noqa: E402
from pydantic import BaseModel, Field  # noqa: E402

from indextts.infer_v2 import IndexTTS2  # noqa: E402

# IndexTTS-2's hardcoded output rate (infer_v2.py:527). Mono int16.
SR = 22050
CHANNELS = 1
BPS = 16

# ── config from env ──────────────────────────────────────────────────
MODEL_DIR = os.environ.get("INDEX_TTS_MODEL_DIR", "/app/checkpoints")
CFG_PATH = os.environ.get("INDEX_TTS_CFG", f"{MODEL_DIR}/config.yaml")
VOICES_DIR = Path(os.environ.get("INDEX_TTS_VOICES_DIR", "/app/voices"))
EMOTIONS_DIR = Path(os.environ.get("INDEX_TTS_EMOTIONS_DIR", "/app/emotions"))
USE_FP16 = os.environ.get("INDEX_TTS_FP16", "1") == "1"
DEVICE = os.environ.get("INDEX_TTS_DEVICE") or None  # "cuda:0", "cuda:1", or None=auto

VOICES_DIR.mkdir(parents=True, exist_ok=True)
EMOTIONS_DIR.mkdir(parents=True, exist_ok=True)

logging.basicConfig(level=os.environ.get("INDEX_TTS_LOG_LEVEL", "INFO"))
log = logging.getLogger("index-tts")

log.info(
    "Loading IndexTTS2 from %s (device=%s, fp16=%s)",
    MODEL_DIR, DEVICE or "auto", USE_FP16,
)
tts = IndexTTS2(
    cfg_path=CFG_PATH,
    model_dir=MODEL_DIR,
    use_fp16=USE_FP16,
    device=DEVICE,
)
log.info("IndexTTS2 ready")

app = FastAPI(title="index-tts", version="0.2.0")


class SpeechRequest(BaseModel):
    model: Optional[str] = "index-tts-2"           # accepted but ignored
    input: str = Field(..., description="Text to synthesize")
    voice: str = Field(..., description="<name>.wav must exist in voices dir")
    response_format: str = Field("wav", description="wav (only)")
    stream: bool = Field(
        False,
        description=(
            "If true, stream the WAV as it generates. Each text segment "
            "(~120 tokens) yields a chunk as soon as IndexTTS-2 finishes "
            "synthesizing it; inter-segment silence is yielded between. "
            "Time-to-first-audio drops dramatically for long inputs. The "
            "WAV header carries placeholder data size (0xFFFFFFFF) so most "
            "players read until EOF."
        ),
    )
    # ── emotion (all optional, mutually exclusive) ──
    emotion_voice: Optional[str] = Field(
        None, description="<name>.wav in emotions dir, used as emotion ref"
    )
    emotion_vector: Optional[List[float]] = Field(
        None,
        description="8 floats: happy, angry, sad, afraid, disgusted, melancholic, surprised, calm",
    )
    emotion_text: Optional[str] = Field(
        None, description="Text describing emotion; QwenEmotion derives vector"
    )
    emotion_alpha: float = Field(1.0, ge=0.0, le=1.0)


def _wav_header(sample_rate: int = SR, channels: int = CHANNELS,
                bits_per_sample: int = BPS) -> bytes:
    """44-byte RIFF/WAVE/PCM header with placeholder data length so the
    payload can be streamed without knowing total samples up front.
    Players that read until EOF (mpv, ffplay, aplay, sox, browsers) handle
    this fine. Strict parsers (some metadata extractors) may complain."""
    byte_rate = sample_rate * channels * bits_per_sample // 8
    block_align = channels * bits_per_sample // 8
    placeholder = 0xFFFFFFFF
    return (
        b"RIFF"
        + placeholder.to_bytes(4, "little")
        + b"WAVE"
        + b"fmt "
        + (16).to_bytes(4, "little")
        + (1).to_bytes(2, "little")                  # PCM
        + channels.to_bytes(2, "little")
        + sample_rate.to_bytes(4, "little")
        + byte_rate.to_bytes(4, "little")
        + block_align.to_bytes(2, "little")
        + bits_per_sample.to_bytes(2, "little")
        + b"data"
        + placeholder.to_bytes(4, "little")
    )


def _chunk_to_pcm_bytes(chunk) -> bytes:
    """Normalize whatever IndexTTS-2 yields (torch tensor, numpy array,
    int16 or float) into raw little-endian int16 PCM bytes."""
    if hasattr(chunk, "cpu"):                        # torch.Tensor
        chunk = chunk.cpu().numpy()
    arr = np.asarray(chunk).reshape(-1)              # flatten to mono samples
    if arr.dtype != np.int16:
        # If the model yields float (-1..1), scale into int16 range.
        if np.issubdtype(arr.dtype, np.floating):
            arr = np.clip(arr * 32767.0, -32768, 32767).astype(np.int16)
        else:
            arr = arr.astype(np.int16)
    return arr.tobytes()


def _resolve(name: str, root: Path) -> Path:
    p = root / f"{name}.wav"
    if not p.is_file():
        raise HTTPException(status_code=404, detail=f"not found: {root.name}/{name}.wav")
    return p


@app.get("/healthz")
def healthz() -> dict:
    return {"status": "ok"}


@app.get("/v1/voices")
def list_voices() -> dict:
    return {
        "voices":   sorted(p.stem for p in VOICES_DIR.glob("*.wav")),
        "emotions": sorted(p.stem for p in EMOTIONS_DIR.glob("*.wav")),
    }


@app.post("/v1/audio/speech")
def synthesize(req: SpeechRequest):
    if req.response_format != "wav":
        raise HTTPException(status_code=400, detail="only response_format=wav is supported")

    spk = str(_resolve(req.voice, VOICES_DIR))

    # First emotion source set wins.
    emo_path = None
    emo_vector = None
    use_emo_text = False
    emo_text = None
    if req.emotion_voice:
        emo_path = str(_resolve(req.emotion_voice, EMOTIONS_DIR))
    elif req.emotion_vector is not None:
        if len(req.emotion_vector) != 8:
            raise HTTPException(status_code=400, detail="emotion_vector must have 8 elements")
        emo_vector = list(req.emotion_vector)
    elif req.emotion_text:
        use_emo_text = True
        emo_text = req.emotion_text

    common_kwargs = dict(
        spk_audio_prompt=spk,
        text=req.input,
        output_path=None,
        emo_audio_prompt=emo_path,
        emo_alpha=req.emotion_alpha,
        emo_vector=emo_vector,
        use_emo_text=use_emo_text,
        emo_text=emo_text,
        verbose=False,
    )

    if not req.stream:
        sr, audio = tts.infer(**common_kwargs)
        buf = io.BytesIO()
        sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
        return Response(content=buf.getvalue(), media_type="audio/wav")

    # Streaming path. tts.infer with stream_return=True is a generator
    # yielding torch audio tensors per text segment plus inter-segment
    # silence. We write a streaming-friendly WAV header up front, then
    # int16 PCM bytes per chunk. Sample rate / channels / bps are fixed
    # by IndexTTS-2 (22050 Hz mono 16-bit).
    def iter_wav() -> Iterable[bytes]:
        yield _wav_header()
        for chunk in tts.infer(stream_return=True, **common_kwargs):
            yield _chunk_to_pcm_bytes(chunk)

    return StreamingResponse(iter_wav(), media_type="audio/wav")