"""PyHarp front end for the original IndexTTS2 inference implementation."""

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import sys
import uuid
from pathlib import Path

import gradio as gr
import soundfile as sf
from pyharp import Channels, ModelCard, SampleRate, Subcategory, build_endpoint

ROOT = Path(__file__).resolve().parent
sys.path.append(str(ROOT / "indextts"))

from indextts.infer_v2 import IndexTTS2  # noqa: E402
from tools.download_files import download_model_from_huggingface  # noqa: E402

model_card = ModelCard(
    name="IndexTTS2",
    description=(
        "Zero-shot speech synthesis with a voice reference and optional emotion "
        "reference, as 22.05 kHz mono speech. Repository: https://github.com/index-tts/index-tts "
        "Paper: https://arxiv.org/abs/2506.21619 Released under the bilibili Model Use "
        "License, which forbids using the model or its output to train or "
        "improve another AI model, except a non-commercial one."
    ),
    author="Siyi Zhou, Yiquan Zhou, Yi He, Xun Zhou, Jinchao Wang, Wei Deng, Jingchen Shu",
    tags=[Subcategory.TEXT_TO_SPEECH, SampleRate(22050), Channels(1), 'voice-cloning',
          'emotion-control'],
)

MODEL_DIR = ROOT / "checkpoints"
OUTPUT_DIR = ROOT / "outputs"
OUTPUT_DIR.mkdir(exist_ok=True)

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is downloaded and loaded there alone
if __name__ != "__main__":
    download_model_from_huggingface(str(MODEL_DIR), str(MODEL_DIR / "hf_cache"))
    tts = IndexTTS2(
        model_dir=str(MODEL_DIR),
        cfg_path=str(MODEL_DIR / "config.yaml"),
        use_fp16=False,
        use_deepspeed=False,
        use_cuda_kernel=False,
    )


def process_fn(text, voice_reference, emotion_reference, emotion_weight):
    # Checked before the GPU is reserved, so a request missing an input spends none of the
    # caller's ZeroGPU allowance
    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter text to synthesize.")
    if not voice_reference:
        raise gr.Error("Please upload a voice reference recording.")
    if gpu_seconds(text, voice_reference, emotion_reference, emotion_weight) > MAX_GPU_SECONDS:
        raise gr.Error(
            f"The text is {len(text)} characters long, more than one request can speak in the "
            f"{MAX_GPU_SECONDS // 60} minutes it may run. Split it."
        )
    for reference in (voice_reference, emotion_reference):
        if reference:
            try:
                info = sf.info(reference)
            except Exception as exc:
                raise gr.Error(f"Could not read the reference recording: {exc}") from exc
            if info.frames == 0:
                raise gr.Error("The reference recording is empty.")

    return speak(text, voice_reference, emotion_reference, emotion_weight)


# The text is spoken a segment at a time, so the cost grows with its length
# PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose
# reservation would exceed that is refused rather than started. A shorter one is left to
# ZeroGPU, which admits it when the caller's own allowance covers it.
MAX_GPU_SECONDS = 900


def gpu_seconds(text, voice_reference, emotion_reference, emotion_weight):
    # Measured at about 12 s for 70 characters, 29 s for 550 and 54 s for 1000, padded
    return 15 + 0.07 * len(text)


@spaces.GPU(duration=gpu_seconds)
def speak(text, voice_reference, emotion_reference, emotion_weight):
    output_path = OUTPUT_DIR / f"harp_{uuid.uuid4().hex}.wav"
    try:
        return tts.infer(
            spk_audio_prompt=voice_reference,
            text=text,
            output_path=str(output_path),
            emo_audio_prompt=emotion_reference or None,
            emo_alpha=float(emotion_weight) * 0.8 if emotion_reference else 1.0,
            max_text_tokens_per_segment=120,
        )
    except Exception as exc:
        raise gr.Error(f"IndexTTS2 synthesis failed: {exc}") from exc


# The processing worker imports this file, so the app must not be built there
if __name__ == "__main__":
    with gr.Blocks(title="IndexTTS2") as demo:
        input_components = [
            gr.Textbox(label="Text to Synthesize", lines=3).set_info(
                "Chinese or English text to synthesize."
            ),
            gr.Audio(type="filepath", label="Voice Reference").set_info(
                "Required voice sample, of which the model uses up to 15 seconds."
            ),
            gr.Audio(type="filepath", label="Emotion Reference")
            .harp_required(False)
            .set_info("Optional. Leave empty to follow the voice reference."),
            gr.Slider(0.0, 1.0, value=0.8, step=0.05, label="Emotion Weight").set_info(
                "Strength of the optional emotion reference."
            ),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Synthesized Speech").set_info(
                "The text spoken in the reference voice."
            ),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)