import tempfile
import gc
import os

import gradio as gr
import soundfile as sf
import torch

from qwen_tts import Qwen3TTSModel

MODEL_ID = "Cartik/Sonexa-1-TTS"

# --- Определяем, доступна ли среда ZeroGPU / CUDA ---
# ZeroGPU (Hugging Face Spaces с декоратором @spaces.GPU) добавляет
# переменную окружения SPACES_ZERO_GPU при реальном запуске на их инфраструктуре.
IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") == "true"
HAS_CUDA = torch.cuda.is_available()

USE_GPU = IS_ZERO_GPU or HAS_CUDA

if USE_GPU:
    DEVICE = "cuda"
    DTYPE = torch.bfloat16
else:
    DEVICE = "cpu"
    DTYPE = torch.float32  # bfloat16 плохо/медленно работает на CPU

print(f"ZeroGPU detected: {IS_ZERO_GPU}")
print(f"CUDA available: {HAS_CUDA}")
print(f"Using device: {DEVICE}, dtype: {DTYPE}")

# --- Условный импорт декоратора spaces.GPU ---
# На обычном CPU-хосте (или локально) пакет spaces может либо отсутствовать,
# либо декоратор просто не нужен — тогда используем no-op заглушку.
if USE_GPU:
    try:
        import spaces

        gpu_decorator = spaces.GPU(duration=60)
    except ImportError:
        print("Пакет 'spaces' не найден, GPU-декоратор отключен.")

        def gpu_decorator(fn):
            return fn
else:
    def gpu_decorator(fn):
        return fn


print("Loading Custom Voice model...")

model = Qwen3TTSModel.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    device_map="auto" if USE_GPU else "cpu",
)

print("Model loaded successfully!")

SUPPORTED_SPEAKERS = [
    "serena",
    "ryan",
    "aiden",
    "dylan",
    "eric",
    "ono_anna",
    "sohee",
    "uncle_fu",
    "vivian",
]


@gpu_decorator
def generate(text, speaker):
    text = (text or "").strip()

    if not text:
        raise gr.Error("Введите текст.")

    try:
        with torch.inference_mode():
            wavs, sr = model.generate_custom_voice(
                text=text,
                speaker=speaker,
                language="Russian",
                max_new_tokens=256,
            )

        tmp = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False,
        )

        sf.write(tmp.name, wavs[0], sr)

        return tmp.name

    except Exception as e:
        raise gr.Error(str(e))

    finally:
        gc.collect()
        if HAS_CUDA:
            torch.cuda.empty_cache()


with gr.Blocks(title="Sonexa TTS") as demo:
    gr.Markdown("# Sonexa TTS (Custom Voice Mode)")

    if USE_GPU:
        gr.Markdown(
            "Генерация выполняется на GPU (ZeroGPU) — обычно занимает несколько секунд."
        )
    else:
        gr.Markdown(
            "⚠️ GPU недоступен — генерация выполняется на CPU и может занимать значительно дольше."
        )

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Текст",
                placeholder="Введите текст...",
                lines=4,
            )

            speaker = gr.Dropdown(
                choices=SUPPORTED_SPEAKERS,
                value="serena",
                label="Выберите голос",
            )

            btn = gr.Button("Озвучить", variant="primary")

        with gr.Column():
            output = gr.Audio(
                label="Результат",
                type="filepath",
            )

    btn.click(
        fn=generate,
        inputs=[input_text, speaker],
        outputs=output,
        api_name="predict",
    )

demo.queue(
    max_size=32,
    default_concurrency_limit=1,
)

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_api=True,
)