"""Thin FastAPI wrapper exposing OmniVoice (k2-fsa/OmniVoice) as an OpenAI-style
TTS for the asset-engine.

Upstream ships only a Gradio demo; we own this wrapper (same pattern as
stacks/index-tts/app.py). Voices are reference WAVs staged in
${OMNIVOICE_VOICES_DIR} (the reused chatterbox /refs/*.wav). A voice-clone prompt
is precomputed once per voice at startup (the loaded Whisper ASR auto-transcribes
each reference) and cached, so per-request latency is just generation.

Endpoints:
  GET  /healthz            -> readiness (200 once model + >=1 voice are loaded)
  GET  /v1/audio/voices    -> {"voices": [<name>, ...]}
  POST /v1/audio/speech    -> {input, voice, response_format=wav} -> audio/wav
"""
import glob
import io
import logging
import os
from pathlib import Path
from typing import Optional

import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel

from omnivoice import OmniVoice, OmniVoiceGenerationConfig

logging.basicConfig(level=os.environ.get("OMNIVOICE_LOG_LEVEL", "INFO"))
log = logging.getLogger("omnivoice-api")

CKPT = os.environ.get("OMNIVOICE_CKPT", "k2-fsa/OmniVoice")
VOICES_DIR = os.environ.get("OMNIVOICE_VOICES_DIR", "/app/voices")
ASR_MODEL = os.environ.get("OMNIVOICE_ASR_MODEL", "openai/whisper-large-v3-turbo")
NUM_STEP = int(os.environ.get("OMNIVOICE_NUM_STEP", "32"))
GUIDANCE = float(os.environ.get("OMNIVOICE_GUIDANCE_SCALE", "2.0"))

app = FastAPI(title="OmniVoice TTS (asset-engine wrapper)")

MODEL: Optional[OmniVoice] = None
PROMPTS: dict = {}        # voice name -> VoiceClonePrompt
SR: int = 24000


class SpeechRequest(BaseModel):
    input: str
    voice: str
    response_format: str = "wav"
    model: Optional[str] = None       # ignored (single model); OpenAI-compat field


@app.on_event("startup")
def _load() -> None:
    global MODEL, SR
    device = "cuda" if torch.cuda.is_available() else "cpu"
    log.info("loading OmniVoice %s on %s (asr=%s)", CKPT, device, ASR_MODEL)
    MODEL = OmniVoice.from_pretrained(
        CKPT, device_map=device, load_asr=True, asr_model_name=ASR_MODEL
    )
    SR = int(getattr(MODEL, "sampling_rate", 24000))

    for wav in sorted(glob.glob(os.path.join(VOICES_DIR, "*.wav"))):
        name = Path(wav).stem
        if name.startswith("_"):
            continue          # skip _*.wav (chatterbox test/deploy artifacts)
        try:
            PROMPTS[name] = MODEL.create_voice_clone_prompt(ref_audio=wav)
            log.info("voice ready: %s", name)
        except Exception as exc:                          # noqa: BLE001
            log.warning("voice %s failed to load: %s", name, exc)
    log.info("%d voices loaded: %s", len(PROMPTS), sorted(PROMPTS))


@app.get("/healthz")
def healthz():
    if MODEL is None or not PROMPTS:
        raise HTTPException(status_code=503, detail="not ready")
    return {"status": "ok", "voices": len(PROMPTS), "sampling_rate": SR}


@app.get("/v1/audio/voices")
def voices():
    return {"voices": sorted(PROMPTS.keys())}


@app.post("/v1/audio/speech")
def speech(req: SpeechRequest):
    if MODEL is None:
        raise HTTPException(status_code=503, detail="model not loaded")
    if not req.input or not req.input.strip():
        raise HTTPException(status_code=400, detail="input is empty")
    if req.response_format not in ("wav", "", None):
        raise HTTPException(status_code=400, detail="only response_format=wav is supported")

    prompt = PROMPTS.get(req.voice)
    if prompt is None:
        raise HTTPException(
            status_code=404,
            detail=f"unknown voice '{req.voice}'; have {sorted(PROMPTS)}",
        )

    gen = OmniVoiceGenerationConfig(num_step=NUM_STEP, guidance_scale=GUIDANCE, denoise=True)
    out = MODEL.generate(text=req.input.strip(), voice_clone_prompt=prompt, generation_config=gen)
    audio = out[0] if isinstance(out, (list, tuple)) else out
    audio = np.asarray(audio, dtype=np.float32)

    buf = io.BytesIO()
    sf.write(buf, audio, SR, format="WAV", subtype="PCM_16")
    return Response(content=buf.getvalue(), media_type="audio/wav")