"""Instant Voice Cloning (F5-TTS / Chatterbox) & Local ZeroGPU Neural TTS Studio."""

import os
import shutil
import tempfile

os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")

import spaces  # Must be imported before torch
import gradio as gr
import numpy as np
import scipy.io.wavfile as wavfile
import torch
from gradio_client import Client, handle_file
from transformers import AutoProcessor, BarkModel

HF_TOKEN = os.environ.get("HF_TOKEN")
LOCAL_TTS_ID = "suno/bark-small"

VOICE_PRESETS = {
    "English Narrator · Studio Male (v2/en_speaker_6)": "v2/en_speaker_6",
    "English Narrator · Warm Female (v2/en_speaker_9)": "v2/en_speaker_9",
    "English Broadcast · Deep Male (v2/en_speaker_0)": "v2/en_speaker_0",
    "English Conversational · Male (v2/en_speaker_3)": "v2/en_speaker_3",
}

print(f"[voice-clone-studio] Preloading {LOCAL_TTS_ID} at module scope for ZeroGPU...", flush=True)
processor = AutoProcessor.from_pretrained(LOCAL_TTS_ID)
bark_model = BarkModel.from_pretrained(LOCAL_TTS_ID, torch_dtype=torch.float16)
bark_model.to("cuda")
SAMPLE_RATE = bark_model.generation_config.sample_rate
print("[voice-clone-studio] Local ZeroGPU TTS ready!", flush=True)


@spaces.GPU(duration=45)
def synthesize_local_tts(text: str, voice_preset_key: str) -> tuple[str, str]:
    """Generate expressive speech locally on ZeroGPU using Bark."""
    preset = VOICE_PRESETS.get(voice_preset_key, "v2/en_speaker_6")
    inputs = processor(text.strip(), voice_preset=preset).to("cuda")
    with torch.no_grad():
        audio_array = bark_model.generate(**inputs)
    audio_np = audio_array.cpu().numpy().squeeze().astype(np.float32)
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        wavfile.write(tmp.name, rate=SAMPLE_RATE, data=audio_np)
        return tmp.name, f"Local ZeroGPU Neural TTS ({preset} · {SAMPLE_RATE} Hz)"


def clone_or_speak(
    text_to_speak: str,
    reference_audio_path: str | None = None,
    reference_transcript: str = "",
    voice_preset_key: str = "English Narrator · Studio Male (v2/en_speaker_6)",
) -> tuple[str, str]:
    """Clone a voice from a 5-10s audio clip (via F5-TTS ZeroGPU) or generate studio neural narration."""
    if not text_to_speak or not text_to_speak.strip():
        raise gr.Error("Please enter the text you want spoken.")

    if reference_audio_path:
        for space_id in ("mrfakename/E2-F5-TTS", "SWivid/F5-TTS"):
            try:
                client = Client(space_id, token=HF_TOKEN, httpx_kwargs={"timeout": 180.0})
                res = client.predict(
                    handle_file(reference_audio_path),
                    (reference_transcript or "").strip(),
                    text_to_speak.strip(),
                    False,
                    api_name="/predict",
                )
                wav_out = res[0] if isinstance(res, (list, tuple)) else res
                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                    shutil.copyfile(wav_out, tmp.name)
                    return tmp.name, f"Zero-Shot Voice Clone via {space_id}"
            except Exception as err:  # noqa: BLE001
                print(f"[voice-clone-studio] {space_id} fallback ({err!r})", flush=True)

    return synthesize_local_tts(text_to_speak, voice_preset_key)


with gr.Blocks(title="Instant Voice Cloning & TTS Studio", fill_width=True) as demo:
    gr.Markdown(
        "# 🎙️ Instant Voice Cloning & Neural TTS Studio\n"
        "Upload a **5–10 second voice clip** for **Zero-Shot Voice Cloning (F5-TTS)**, or leave the reference audio blank "
        "to generate studio narration on **ZeroGPU (`zero-a10g`)** at **$0/hour**."
    )
    with gr.Row():
        with gr.Column(scale=1):
            text_box = gr.Textbox(
                label="Script / Text to Speak",
                lines=5,
                placeholder="Type the narration, dialogue, or voiceover script here...",
            )
            ref_audio = gr.Audio(
                label="Reference Voice Audio (Optional · 5–10s clip for Voice Cloning)",
                type="filepath",
            )
            ref_text = gr.Textbox(
                label="Reference Audio Transcript (Optional · improves F5-TTS voice clone accuracy)",
                lines=2,
                placeholder="Leave blank for automatic transcription, or type what was said in the reference clip...",
            )
            preset_dropdown = gr.Dropdown(
                choices=list(VOICE_PRESETS),
                value="English Narrator · Studio Male (v2/en_speaker_6)",
                label="Studio Voice Preset (Used when no reference audio is uploaded)",
            )
            speak_btn = gr.Button("🎙️ Generate Voiceover", variant="primary")
        with gr.Column(scale=1):
            audio_out = gr.Audio(label="Generated Voiceover (WAV)", type="filepath", autoplay=True)
            status_out = gr.Textbox(label="Engine Details", interactive=False)

    speak_btn.click(
        fn=clone_or_speak,
        inputs=[text_box, ref_audio, ref_text, preset_dropdown],
        outputs=[audio_out, status_out],
        api_name="synthesize_voice",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), mcp_server=True)