import os
import gradio as gr
import base64
from Zonos_main.tts import ZonosTTS

# Force CPU mode to avoid CUDA issues
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

def load_model():
    try:
        tts = ZonosTTS()
        print("✅ TTS model loaded successfully")
        return tts
    except Exception as e:
        print(f"❌ Model loading failed: {str(e)}")
        raise RuntimeError("Could not initialize TTS model") from e

tts = load_model()

def generate_audio(text):
    try:
        if not text.strip():
            raise ValueError("Input text cannot be empty")
        
        # Generate audio (must return raw bytes)
        audio_bytes = tts.generate(text)
        
        if not isinstance(audio_bytes, bytes):
            raise TypeError("Model must return bytes")
            
        if len(audio_bytes) < 100:  # Minimum valid audio size
            raise ValueError("Generated audio too small")
            
        return {
            "name": "output.wav",
            "data": f"data:audio/wav;base64,{base64.b64encode(audio_bytes).decode()}",
            "size": len(audio_bytes)
        }
        
    except Exception as e:
        print(f"❌ Generation error: {str(e)}")
        raise gr.Error(f"Audio generation failed: {str(e)}") from e

# Robust Gradio interface
demo = gr.Interface(
    fn=generate_audio,
    inputs=gr.Textbox(label="Input Text", placeholder="Type something..."),
    outputs=gr.Audio(label="Output", type="filepath"),
    title="Zonos TTS",
    allow_flagging="never"
)

# Launch with error handling
try:
    demo.launch(show_error=True)
except Exception as e:
    print(f"❌ Gradio launch failed: {str(e)}")
    raise