"""GEPARD — text-to-speech inference Space (ZeroGPU), gradio.Server backend.

Startup order is load-bearing:
  1. ``create_env.setup_dependencies()`` installs nemo-toolkit and re-pins
     transformers BEFORE any ML import — the Space image ships gradio 6.20
     which pins huggingface-hub>=1.2, and NeMo's transformers<=4.52 would
     pull hub<1.0 at build time. create_env.py upgrades hub first, then
     installs nemo with --no-deps, then re-pins transformers==5.3.0.
  2. ``spaces`` is imported before torch so ZeroGPU can patch CUDA calls.
  3. The engine (model + codec + speakers) is built ONCE at module level;
     ZeroGPU replays the recorded ``.to("cuda")`` moves when a GPU is
     attached, so no per-request reloading happens.

The UI is a custom HTML/CSS/JS frontend served at ``/``; ``@app.api()``
exposes the synthesis pipeline through Gradio's queue so the JS client
(and ``gradio_client``) can call it without colliding on the GPU.
"""

from create_env import setup_dependencies

setup_dependencies()

try:  # ZeroGPU runtime (present on HF Spaces)
    import spaces  # noqa: E402

    _gpu_decorator = spaces.GPU
except ImportError:  # local run — no-op stand-in
    def _gpu_decorator(*args, **kwargs):
        if args and callable(args[0]):
            return args[0]

        def _wrap(fn):
            return fn

        return _wrap

import os  # noqa: E402
from pathlib import Path  # noqa: E402

import numpy as np  # noqa: E402
import scipy.io.wavfile as wavfile  # noqa: E402

from fastapi.responses import HTMLResponse  # noqa: E402

from gradio import Server  # noqa: E402
from gradio.data_classes import FileData  # noqa: E402

from gepard_inference.engine import AppConfig, GenerationParams, GepardEngine  # noqa: E402
from interface import MODE_PRESET  # noqa: E402

CONFIG_PATH = Path(__file__).parent / "config.yaml"

config = AppConfig.from_yaml(CONFIG_PATH)
engine = GepardEngine(config).load()

# Where synthesized WAVs land — same dir as the script so Spaces' persistent
# volume picks them up across restarts.
OUT_DIR = Path(__file__).parent / "_out"
OUT_DIR.mkdir(exist_ok=True)


def _save_wav(sample_rate: int, wave: np.ndarray) -> str:
    """Write a float waveform to a 16-bit PCM WAV; return the absolute path."""
    import uuid

    # Clip & convert to int16 PCM (browser audio tag handles 16-bit PCM natively).
    wave = np.asarray(wave, dtype=np.float32)
    peak = float(np.max(np.abs(wave))) if wave.size else 0.0
    if peak > 1.0:
        wave = wave / peak
    pcm = (wave * 32767.0).clip(-32768, 32767).astype(np.int16)
    path = OUT_DIR / f"{uuid.uuid4().hex}.wav"
    wavfile.write(path, sample_rate, pcm)
    return str(path)


app = Server()


@app.api()
@_gpu_decorator(duration=config.gpu_duration)
def synthesize(
    mode: str,
    speaker: str,
    ref_audio: str,
    text: str,
    temperature: float,
    top_k: float,
    max_frames: float,
    repetition_penalty: float,
    repetition_window: float,
) -> FileData:
    """Generate speech from text + a preset speaker or a recorded reference.

    Runs inside the ZeroGPU context — both the reference encoding (clone
    mode) and the autoregressive generation share one GPU session.

    Args match the current Gradio Blocks signature exactly so the rest of
    the codebase (config defaults, examples) stays a drop-in.
    """
    text = (text or "").strip()
    if not text:
        raise ValueError("Please enter some text to synthesize.")

    if mode == MODE_PRESET:
        if not speaker:
            raise ValueError("Please choose a preset speaker.")
        ref_codes = engine.speakers.codes(speaker)
    else:
        if not ref_audio:
            raise ValueError("Please record or upload a reference clip.")
        ref_codes = engine.encode_reference(ref_audio)

    params = GenerationParams(
        temperature=temperature,
        top_k=int(top_k),
        cfg_scale=config.defaults.cfg_scale,
        cfg_frames=config.defaults.cfg_frames,
        stop_threshold=config.defaults.stop_threshold,
        max_frames=int(max_frames),
        repetition_penalty=repetition_penalty,
        repetition_window=int(repetition_window),
    )
    sample_rate, wave = engine.synthesize(text, ref_codes, params)
    return FileData(path=_save_wav(sample_rate, wave))


@app.get("/")
async def homepage():
    """Serve the custom HTML/CSS/JS frontend.

    Server is a FastAPI app, so plain ``@app.get`` routes coexist naturally
    with ``@app.api()`` routes. The frontend lives in ``index.html`` next
    to this file. Returning an ``HTMLResponse`` (instead of a plain string)
    sets ``Content-Type: text/html`` so the browser actually renders it.
    """
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


@app.get("/config.json")
async def public_config():
    """Static config exposed to the frontend (speaker list, defaults, examples)."""
    return {
        "speakers": [
            {"key": k, "label": config.speaker_labels.get(k, k)}
            for k in engine.speakers.names
        ],
        "defaults": {
            "temperature": config.defaults.temperature,
            "top_k": config.defaults.top_k,
            "max_frames": config.defaults.max_frames,
            "repetition_penalty": config.defaults.repetition_penalty,
            "repetition_window": config.defaults.repetition_window,
        },
        "limits": {
            "min_max_frames": 43,
            "max_max_frames": 1075,
            "min_repetition_window": 0,
            "max_repetition_window": 128,
            "min_repetition_penalty": 1.0,
            "max_repetition_penalty": 1.5,
            "min_temperature": 0.05,
            "max_temperature": 1.0,
            "min_top_k": 0,
            "max_top_k": 200,
        },
        "examples": [
            {"speaker": ex.speaker, "label": ex.label, "text": ex.text}
            for ex in config.examples
        ],
        "modes": {"preset": MODE_PRESET, "clone": "Clone a voice"},
    }


if __name__ == "__main__":
    app.launch(show_error=True)