#!/usr/bin/env python3
"""SOTAKA — Speak (VoxCPM2) + Song Convert (Seed-VC) on Hugging Face Spaces."""

from __future__ import annotations

import os

os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")

from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
import numpy as np
import soundfile as sf

try:
    import spaces
except ImportError:  # local / non-ZeroGPU

    class spaces:  # type: ignore
        @staticmethod
        def GPU(*_args, **_kwargs):
            def deco(fn):
                return fn

            return deco


MODEL_ID = os.environ.get("VOXCPM_MODEL", "openbmb/VoxCPM2")
OUT_DIR = Path("outputs")
OUT_DIR.mkdir(exist_ok=True)

_model = None


def get_model():
    global _model
    if _model is None:
        from voxcpm import VoxCPM

        _model = VoxCPM.from_pretrained(MODEL_ID, load_denoiser=False)
    return _model


def _build_text(text: str, instruct: Optional[str]) -> str:
    if instruct:
        return f"({instruct}){text}"
    return text


@spaces.GPU(duration=90)
def synthesize(
    text: str,
    instruct: str,
    ref_audio: Optional[str],
    ref_text: str,
    timesteps: int,
    cfg: float,
) -> Tuple[Optional[str], str]:
    text = (text or "").strip()
    if not text:
        return None, "Please enter text."

    instruct = (instruct or "").strip() or None
    ref_text = (ref_text or "").strip() or None

    try:
        model = get_model()
    except Exception as e:
        return None, f"Failed to load model: {e}"

    kwargs = {
        "text": _build_text(text, instruct),
        "inference_timesteps": int(timesteps),
        "cfg_value": float(cfg),
    }
    if ref_audio:
        kwargs["reference_wav_path"] = ref_audio
        if ref_text:
            kwargs["prompt_wav_path"] = ref_audio
            kwargs["prompt_text"] = ref_text

    try:
        wav = model.generate(**kwargs)
        if wav is None or (hasattr(wav, "size") and wav.size == 0):
            return None, "No audio generated."
        wav = np.asarray(wav, dtype=np.float32).reshape(-1)
        sample_rate = int(getattr(model.tts_model, "sample_rate", 48000))
        out = OUT_DIR / "SOTAKA_Voice.wav"
        sf.write(out, wav, sample_rate)
        dur = len(wav) / sample_rate
        return str(out), f"Model: {MODEL_ID}\nduration={dur:.2f}s | sr={sample_rate}"
    except Exception as e:
        return None, f"Generation error: {e}"


# Lazy import so Speak still boots if Song deps fail to resolve at import time.
from song_convert import convert_song  # noqa: E402


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="SOTAKA") as demo:
        gr.Markdown(
            """
# SOTAKA
**Speak** = text → voice (VoxCPM2) · **Song Convert** = keep melody, swap timbre (Seed-VC)
            """
        )
        with gr.Tabs():
            with gr.Tab("Speak"):
                gr.Markdown(
                    """
Type Khmer (or other) text. Optional: voice design instruction, or a short reference clip to clone.
                    """
                )
                text = gr.Textbox(
                    label="Text",
                    lines=4,
                    value="សួស្តី។ ខ្ញុំសប្បាយចិត្តដែលបានជួបអ្នក។",
                )
                with gr.Row():
                    instruct = gr.Textbox(
                        label="Voice design (optional)",
                        placeholder="A calm adult Khmer male voice, clear and natural",
                        lines=2,
                    )
                    ref_audio = gr.Audio(
                        label="Reference audio (optional clone)",
                        type="filepath",
                    )
                ref_text = gr.Textbox(
                    label="Reference transcript (optional)",
                    lines=2,
                )
                with gr.Row():
                    timesteps = gr.Slider(4, 16, value=5, step=1, label="Inference steps")
                    cfg = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="CFG")
                btn_speak = gr.Button("Generate", variant="primary")
                speak_out = gr.Audio(label="Output", type="filepath")
                speak_log = gr.Textbox(label="Log", lines=4)
                btn_speak.click(
                    synthesize,
                    inputs=[text, instruct, ref_audio, ref_text, timesteps, cfg],
                    outputs=[speak_out, speak_log],
                )

            with gr.Tab("Song Convert"):
                gr.Markdown(
                    """
Keep **melody & timing**, swap to **your voice**.

Tips: source = **vocals only** · max **60s** · reference = **8–20s** clear speech · use **ZeroGPU** or **T4**
                    """
                )
                with gr.Row():
                    source = gr.Audio(label="Source vocal / song", type="filepath")
                    voice_ref = gr.Audio(label="Your voice reference", type="filepath")
                with gr.Row():
                    steps = gr.Slider(10, 80, value=30, step=1, label="Diffusion steps")
                    pitch = gr.Slider(-12, 12, value=0, step=1, label="Pitch shift (semitones)")
                with gr.Row():
                    length = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Length adjust")
                    song_cfg = gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="CFG rate")
                btn_song = gr.Button("Convert to my voice", variant="primary")
                song_out = gr.Audio(label="Output", type="filepath")
                song_log = gr.Textbox(label="Log", lines=5)
                btn_song.click(
                    convert_song,
                    inputs=[source, voice_ref, steps, pitch, length, song_cfg],
                    outputs=[song_out, song_log],
                )
    return demo


demo = build_ui()
demo.queue(max_size=8)

if __name__ == "__main__":
    print(f"Starting UI. Speak model: {MODEL_ID}")
    demo.launch(ssr_mode=False)