# Qwen3-TTS — Voice Clone only (Base 1.7B)
import os
import gradio as gr
import numpy as np
import torch
from huggingface_hub import snapshot_download
from qwen_tts import Qwen3TTSModel

MODEL_SIZES = ["0.6B", "1.7B"]
LANGUAGES = ["Auto", "Chinese", "English", "Japanese", "Korean", "French", "German", "Spanish", "Portuguese", "Russian", "Italian"]


def get_model_path(model_type: str, model_size: str) -> str:
    return snapshot_download(f"Qwen/Qwen3-TTS-12Hz-{model_size}-{model_type}")


print("Loading Base 1.7B model...")
base_model_1_7b = Qwen3TTSModel.from_pretrained(
    get_model_path("Base", "1.7B"),
    device_map="cuda",
    dtype=torch.float16,
)
print("Model loaded successfully!")

BASE_MODELS = {
    "1.7B": base_model_1_7b,
}


def _normalize_audio(wav, eps=1e-12, clip=True):
    x = np.asarray(wav)
    if np.issubdtype(x.dtype, np.integer):
        info = np.iinfo(x.dtype)
        if info.min < 0:
            y = x.astype(np.float32) / max(abs(info.min), info.max)
        else:
            mid = (info.max + 1) / 2.0
            y = (x.astype(np.float32) - mid) / mid
    elif np.issubdtype(x.dtype, np.floating):
        y = x.astype(np.float32)
        m = np.max(np.abs(y)) if y.size else 0.0
        if m > 1.0 + 1e-6:
            y = y / (m + eps)
    else:
        raise TypeError(f"Unsupported dtype: {x.dtype}")
    if clip:
        y = np.clip(y, -1.0, 1.0)
    if y.ndim > 1:
        y = np.mean(y, axis=-1).astype(np.float32)
    return y


def _audio_to_tuple(audio):
    if audio is None:
        return None
    if isinstance(audio, tuple) and len(audio) == 2 and isinstance(audio[0], int):
        sr, wav = audio
        wav = _normalize_audio(wav)
        return wav, int(sr)
    if isinstance(audio, dict) and "sampling_rate" in audio and "data" in audio:
        sr = int(audio["sampling_rate"])
        wav = _normalize_audio(audio["data"])
        return wav, sr
    return None


def generate_voice_clone(ref_audio, ref_text, target_text, language, use_xvector_only, model_size, progress=gr.Progress(track_tqdm=True)):
    if not target_text or not target_text.strip():
        return None, "Error: Target text is required."
    audio_tuple = _audio_to_tuple(ref_audio)
    if audio_tuple is None:
        return None, "Error: Reference audio is required."
    if not use_xvector_only and (not ref_text or not ref_text.strip()):
        return None, "Error: Reference text is required when 'Use x-vector only' is not enabled."
    try:
        tts = BASE_MODELS.get(model_size, base_model_1_7b)
        wavs, sr = tts.generate_voice_clone(
            text=target_text.strip(),
            language=language,
            ref_audio=audio_tuple,
            ref_text=ref_text.strip() if ref_text else None,
            x_vector_only_mode=use_xvector_only,
            max_new_tokens=2048,
        )
        return (sr, wavs[0]), "Voice clone generation completed successfully!"
    except Exception as e:
        return None, f"Error: {type(e).__name__}: {e}"


def build_ui():
    with gr.Blocks(title="Worder TTS") as demo:
        gr.Markdown("# Worder Voice Clone (Base 1.7B)")
        with gr.Row():
            with gr.Column():
                clone_ref_audio = gr.Audio(label="Reference Audio", type="numpy")
                clone_ref_text = gr.Textbox(label="Reference Text", lines=2)
                clone_xvector = gr.Checkbox(label="Use x-vector only", value=False)
            with gr.Column():
                clone_target_text = gr.Textbox(label="Target Text", lines=4)
                with gr.Row():
                    clone_language = gr.Dropdown(label="Language", choices=LANGUAGES, value="Auto")
                    clone_model_size = gr.Dropdown(label="Model Size", choices=["1.7B"], value="1.7B")
                clone_btn = gr.Button("Generate", variant="primary")
        clone_audio_out = gr.Audio(label="Generated Audio", type="numpy")
        clone_status = gr.Textbox(label="Status", lines=2, interactive=False)
        clone_btn.click(
            generate_voice_clone,
            inputs=[clone_ref_audio, clone_ref_text, clone_target_text, clone_language, clone_xvector, clone_model_size],
            outputs=[clone_audio_out, clone_status],
        )
    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.launch()