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

from __future__ import annotations

import sys
import uuid
from pathlib import Path

import gradio as gr
import spaces
from pyharp import ModelCard, 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_DIR = ROOT / "checkpoints"
OUTPUT_DIR = ROOT / "outputs"
OUTPUT_DIR.mkdir(exist_ok=True)

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 synthesize(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.")

    return speak(text, voice_reference, emotion_reference, emotion_weight)


@spaces.GPU
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


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; 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=ModelCard(
            name="IndexTTS2",
            description=(
                "Zero-shot speech synthesis with a voice reference and optional emotion "
                "reference. 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, 'voice-cloning', 'emotion-control'],
        ),
        input_components=input_components,
        output_components=output_components,
        process_fn=synthesize,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)