import spaces
import os
import sys
import tempfile
import numpy as np
import gradio as gr
import soundfile as sf

# Add NeuTTS-Air to path
sys.path.append("neutts-air")

# Global variables for lazy loading
kokoro_pipe = None
neutts_model = None

# NeuTTS-Air configuration
SAMPLES_PATH = os.path.join(os.getcwd(), "neutts-air", "samples")
DEFAULT_REF_TEXT = "So I'm live on radio. And I say, well, my dear friend James here clearly, and the whole room just froze. Turns out I'd completely misspoken and mentioned our other friend."
DEFAULT_REF_PATH = os.path.join(SAMPLES_PATH, "dave.wav")
DEFAULT_GEN_TEXT = "My name is Dave, and um, I'm from London."

# ------------------------------------------------------------------
# 1. Lazy loaders
# ------------------------------------------------------------------
def load_kokoro():
    global kokoro_pipe
    if kokoro_pipe is None:
        from kokoro import KPipeline
        kokoro_pipe = KPipeline(lang_code='a')
    return kokoro_pipe

def load_neutts():
    global neutts_model
    if neutts_model is None:
        from neuttsair.neutts import NeuTTSAir
        neutts_model = NeuTTSAir(
            backbone_repo="neuphonic/neutts-air",
            backbone_device="cuda",
            codec_repo="neuphonic/neucodec",
            codec_device="cuda"
        )
    return neutts_model

# ------------------------------------------------------------------
# 2. Kokoro TTS inference
# ------------------------------------------------------------------
@spaces.GPU()
def kokoro_infer(text, voice, speed):
    if not text.strip():
        raise gr.Error("Please enter some text.")
    
    pipe = load_kokoro()
    generator = pipe(text, voice=voice, speed=speed)
    for gs, ps, audio in generator:
        # Save to temporary file
        fd, tmp = tempfile.mkstemp(suffix='.wav')
        os.close(fd)
        sf.write(tmp, audio, 24000)
        return tmp
    raise RuntimeError("Kokoro generation failed")

# ------------------------------------------------------------------
# 3. NeuTTS-Air inference
# ------------------------------------------------------------------
@spaces.GPU()
def neutts_infer(ref_text: str, ref_audio_path: str, gen_text: str) -> tuple:
    """
    Generates speech using NeuTTS-Air given a reference audio and text, and new text to synthesize.
    """
    if not gen_text.strip():
        raise gr.Error("Please enter text to generate.")
    if not ref_audio_path:
        raise gr.Error("Please provide reference audio.")
    if not ref_text.strip():
        raise gr.Error("Please provide reference text.")
    
    gr.Info("Starting inference request!")
    gr.Info("Encoding reference...")
    
    tts = load_neutts()
    ref_codes = tts.encode_reference(ref_audio_path)
    
    gr.Info(f"Generating audio for input text: {gen_text}")
    wav = tts.infer(gen_text, ref_codes, ref_text)
    
    return (24_000, wav)

# ------------------------------------------------------------------
# 4. Gradio UI with model selection
# ------------------------------------------------------------------
css = """footer {visibility: hidden}"""

with gr.Blocks(css=css, title="Text2Audio - Kokoro & NeuTTS-Air") as demo:
    gr.Markdown("# 🎙️ Text-to-Audio Generation")
    gr.Markdown("Choose between **Kokoro TTS** (fast English TTS) or **NeuTTS-Air** (voice cloning with reference audio)")
    
    # Model selection
    model_choice = gr.Radio(
        choices=["Kokoro TTS", "NeuTTS-Air"],
        value="Kokoro TTS",
        label="Select TTS Engine",
        interactive=True
    )
    
    # Kokoro TTS Interface
    with gr.Group(visible=True) as kokoro_group:
        gr.Markdown("### 🎧 Kokoro TTS Settings")
        with gr.Row():
            with gr.Column():
                kokoro_voice = gr.Dropdown(
                    label="Voice",
                    choices=['af_heart', 'af_sky', 'af_mist', 'af_dusk'],
                    value='af_heart'
                )
                kokoro_speed = gr.Slider(0.5, 2.0, 1.0, step=0.1, label="Speed")
            
            with gr.Column(scale=3):
                kokoro_text = gr.Textbox(
                    label="Text to speak",
                    placeholder="Type or paste text here…",
                    lines=6,
                    max_lines=12
                )
                kokoro_btn = gr.Button("🎧 Synthesise with Kokoro", variant="primary")
                kokoro_audio_out = gr.Audio(label="Generated speech", type="filepath")
        
        gr.Markdown("**Kokoro** – fast, high-quality English TTS. Audio is returned as 24 kHz WAV.")
    
    # NeuTTS-Air Interface
    with gr.Group(visible=False) as neutts_group:
        gr.Markdown("### ☁️ NeuTTS-Air Settings")
        neutts_ref_text = gr.Textbox(
            label="Reference Text",
            value=DEFAULT_REF_TEXT,
            lines=3
        )
        neutts_ref_audio = gr.Audio(
            type="filepath",
            label="Reference Audio",
            value=DEFAULT_REF_PATH if os.path.exists(DEFAULT_REF_PATH) else None
        )
        neutts_gen_text = gr.Textbox(
            label="Text to Generate",
            value=DEFAULT_GEN_TEXT,
            lines=3
        )
        neutts_btn = gr.Button("☁️ Generate with NeuTTS-Air", variant="primary")
        neutts_audio_out = gr.Audio(label="Generated Speech", type="numpy")
        
        gr.Markdown("**NeuTTS-Air** – Upload a reference audio sample, provide the reference text, and enter new text to synthesize.")
    
    # Event handlers
    def toggle_interface(choice):
        if choice == "Kokoro TTS":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)
    
    model_choice.change(
        fn=toggle_interface,
        inputs=[model_choice],
        outputs=[kokoro_group, neutts_group]
    )
    
    kokoro_btn.click(
        kokoro_infer,
        inputs=[kokoro_text, kokoro_voice, kokoro_speed],
        outputs=kokoro_audio_out
    )
    
    neutts_btn.click(
        neutts_infer,
        inputs=[neutts_ref_text, neutts_ref_audio, neutts_gen_text],
        outputs=neutts_audio_out
    )

if __name__ == "__main__":
    demo.launch(allowed_paths=[SAMPLES_PATH] if os.path.exists(SAMPLES_PATH) else None, mcp_server=True, inbrowser=True)