import os
from pathlib import Path

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import gradio as gr
import torch
from huggingface_hub import snapshot_download

from m5_infer import GenerationOptions, M5TTS
from tn_frontend import get_normalizer


ROOT = Path(__file__).resolve().parent
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DEFAULT_VOICE = ROOT / "assets" / "aslp_default_voice.wav"
DEFAULT_PROMPT = "猪携带该疾病，然后通过蚊子传染给人类。"
DEFAULT_TARGET = (
    "欢迎来到ASLP实验室。我们专注于语音、语言与智能交互技术研究，"
    "期待用声音连接更广阔的世界。"
)
QWEN_ROOT = Path(snapshot_download(
    repo_id="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
    revision="fd4b254389122332181a7c3db7f27e918eec64e3",
    allow_patterns=["speech_tokenizer/*"],
))
get_normalizer("zh")
get_normalizer("en")
tts = M5TTS(
    ROOT / "model" / "model.safetensors",
    QWEN_ROOT / "speech_tokenizer",
    device=DEVICE,
    runtime_root=ROOT,
    expected_sha256="6c1ac8578aafb07d444e284c76d16d6529e87e1971189c76ff1795bf1252d623",
)


def estimate_duration(_audio, _prompt, target, *_args, **_kwargs):
    return min(120, max(30, 20 + len(target or "") // 2))


@spaces.GPU(size="large", duration=estimate_duration)
def synthesize(
    reference_audio,
    prompt_text,
    target_text,
    language,
    accent,
    temperature,
    top_k,
    seed,
):
    if not reference_audio:
        raise gr.Error("Please upload or record reference audio.")
    prompt_text = (prompt_text or "").strip()
    target_text = (target_text or "").strip()
    if not prompt_text:
        raise gr.Error("Reference transcript is required.")
    if not target_text:
        raise gr.Error("Target text is required.")
    if len(target_text) > 250:
        raise gr.Error("Target text is limited to 250 characters per request.")

    accent = None if accent == "None" else accent
    options = GenerationOptions(
        seed=int(seed),
        temperature=float(temperature),
        top_k=int(top_k),
        text_token_limit=96,
        max_audio_ms=30_000,
        max_seq_len=2048,
        allow_truncated=False,
    )
    result = tts.synthesize(
        target_text,
        reference_audio,
        prompt_text,
        language=language,
        reference_language=language,
        accent=accent,
        speaker=None,
        caption=None,
        utt="space",
        tn="wetext",
        options=options,
    )
    info = result["metadata"]
    return (result["sample_rate"], result["audio"]), {
        "sample_rate": result["sample_rate"],
        "duration_seconds": round(info["duration"], 3),
        "chunks": info["chunks"],
        "target_tn": info["prepared"]["target"],
        "reference_tn": info["prepared"]["reference"],
        "conditioning": {
            "language": info["language"],
            "accent": info["accent"],
            "speaker": info["speaker"],
        },
    }


with gr.Blocks(title="M5-TTS") as demo:
    gr.Markdown(
        "# M5-TTS\n"
        "Upload a reference recording and its exact transcript, then synthesize "
        "new speech with the same voice."
    )
    with gr.Row():
        with gr.Column():
            reference_audio = gr.Audio(
                sources=["upload", "microphone"], type="filepath",
                value=str(DEFAULT_VOICE), label="Reference audio",
            )
            prompt_text = gr.Textbox(
                value=DEFAULT_PROMPT, label="Reference transcript", lines=3,
                info="Update this transcript when uploading another voice.",
            )
            target_text = gr.Textbox(
                value=DEFAULT_TARGET, label="Target text", lines=5, max_length=250,
                placeholder="Text to synthesize (maximum 250 characters)",
            )
        with gr.Column():
            language = gr.Dropdown(
                [
                    ("Chinese", "zh"), ("English", "en"), ("German", "de"),
                    ("French", "fr"), ("Japanese", "ja"), ("Korean", "ko"),
                ],
                value="zh", label="Language",
            )
            accent = gr.Dropdown(
                ["None", "yue", "chuan"], value="None", label="Accent",
            )
            temperature = gr.Slider(
                0.1, 1.5, value=0.9, step=0.05, label="Temperature",
            )
            top_k = gr.Slider(1, 50, value=10, step=1, label="Top-k")
            seed = gr.Number(value=42, precision=0, label="Seed")
    generate_button = gr.Button("Generate", variant="primary")
    output_audio = gr.Audio(label="Generated audio")
    metadata = gr.JSON(label="Generation metadata")
    generate_button.click(
        fn=synthesize,
        inputs=[
            reference_audio, prompt_text, target_text, language, accent,
            temperature, top_k, seed,
        ],
        outputs=[output_audio, metadata],
        api_name="synthesize",
        concurrency_limit=1,
    )

demo.queue(default_concurrency_limit=1).launch()