import json
import os

import gradio as gr
import numpy as np
import spaces
import torch
from qwen_tts import Qwen3TTSModel

# 0.6B matches the voices' origin; switch to Qwen/Qwen3-TTS-12Hz-1.7B-Base for
# higher quality at ~2.5x the GPU time per clip.
MODEL_ID = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
MAX_CHARS = 1500

VOICES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "voices")
with open(os.path.join(VOICES_DIR, "transcripts.json"), encoding="utf-8") as f:
    TRANSCRIPTS = json.load(f)
VOICES = sorted(TRANSCRIPTS.keys())
LANGUAGES = ["English", "Chinese", "Japanese", "Korean", "German",
             "French", "Russian", "Portuguese", "Spanish", "Italian"]

model = Qwen3TTSModel.from_pretrained(
    MODEL_ID,
    device_map="cuda",
    dtype=torch.bfloat16,
)

_prompt_cache = {}


def _get_voice_prompt(voice: str):
    if voice not in _prompt_cache:
        _prompt_cache[voice] = model.create_voice_clone_prompt(
            ref_audio=os.path.join(VOICES_DIR, f"{voice}.wav"),
            ref_text=TRANSCRIPTS[voice],
        )
    return _prompt_cache[voice]


@spaces.GPU(duration=90)
def tts(text: str, voice: str, language: str):
    text = (text or "").strip()
    if not text:
        raise gr.Error("Enter some text to speak.")
    if len(text) > MAX_CHARS:
        raise gr.Error(f"Text too long ({len(text)} chars, max {MAX_CHARS}).")
    if voice not in TRANSCRIPTS:
        raise gr.Error(f"Unknown voice '{voice}'. Available: {', '.join(VOICES)}")

    wavs, sr = model.generate_voice_clone(
        text=text,
        language=language,
        voice_clone_prompt=_get_voice_prompt(voice),
    )
    audio = np.asarray(wavs[0], dtype=np.float32)
    return sr, audio


demo = gr.Interface(
    fn=tts,
    inputs=[
        gr.Textbox(label="Text", lines=4, placeholder="What should the voice say?"),
        gr.Dropdown(VOICES, value=VOICES[0], label="Voice"),
        gr.Dropdown(LANGUAGES, value="English", label="Language"),
    ],
    outputs=gr.Audio(label="Generated speech"),
    title="EsfandTTS — Qwen3-TTS voice clone",
    description="Cloned voices via Qwen3-TTS 0.6B Base. Also callable as an API (see the 'Use via API' link below).",
    flagging_mode="never",
)

demo.launch()