import gradio as gr
import numpy as np
from kittentts import KittenTTS

# 1. Initialize the KittenTTS model.
# This will download the model from Hugging Face on the first run.
print("Loading KittenTTS model...")
try:
    tts_model = KittenTTS("KittenML/kitten-tts-nano-0.1")
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    # You might want to handle this more gracefully
    exit()


# 2. Get the list of available voices directly from the model instance.
AVAILABLE_VOICES = tts_model.available_voices
DEFAULT_VOICE = "expr-voice-5-m" if "expr-voice-5-m" in AVAILABLE_VOICES else AVAILABLE_VOICES[0]

# 3. Define the core function that Gradio will call.
# This function now accepts 'voice' and 'speed' as arguments.
def synthesize_speech(text, voice, speed):
    """
    Generates audio using the selected text, voice, and speed.
    """
    # Handle empty input gracefully
    if not text.strip():
        # Return a silent, empty audio clip
        return (24000, np.zeros(0, dtype=np.int16))

    # Call the model's generate method with all the parameters
    audio_data = tts_model.generate(text, voice=voice, speed=speed)
    
    # Return the audio in the format Gradio expects: (sample_rate, numpy_array)
    return (24000, audio_data)

# 4. Create the Gradio UI with the new controls.
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🐱 Enhanced KittenTTS UI
        A user-friendly interface for the KittenTTS text-to-speech model.
        Select a voice, adjust the speed, and type your text to generate audio.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            text_input = gr.Textbox(
                lines=5,
                label="Input Text",
                placeholder="Type something here..."
            )
        
        with gr.Column(scale=1):
            voice_dropdown = gr.Dropdown(
                choices=AVAILABLE_VOICES,
                value=DEFAULT_VOICE,
                label="Voice Selection"
            )
            speed_slider = gr.Slider(
                minimum=0.5,
                maximum=2.0,
                step=0.1,
                value=1.0,
                label="Speech Speed"
            )
            generate_button = gr.Button("Generate Audio", variant="primary")
            
    audio_output = gr.Audio(label="Generated Speech", autoplay=True)

    gr.Examples(
        examples=[
            ["This is an example of a female voice.", "expr-voice-5-f", 1.0],
            ["This is an example of a male voice, speaking a bit faster.", "expr-voice-5-m", 1.2],
            ["The speed can also be slowed down for clarity.", "expr-voice-4-f", 0.8],
        ],
        inputs=[text_input, voice_dropdown, speed_slider]
    )

    # Connect the UI components to the function
    generate_button.click(
        fn=synthesize_speech,
        inputs=[text_input, voice_dropdown, speed_slider],
        outputs=audio_output
    )

# Launch the Gradio app
demo.launch()