from kokoro import KPipeline
import torch
import soundfile as sf
import numpy as np
import spaces

# Initialize the Kokoro TTS pipeline
# Using 'a' as the language code for auto-detection
pipeline = KPipeline(lang_code='a')

# Default voice - you can change this to any of the available voices
# Some options include: 'af_heart', 'en_us_female', etc.
DEFAULT_VOICE = 'af_heart'

@spaces.GPU(duration_s=20)  # TTS typically takes less time than other operations
def synthesize_speech(text, output_path="output.wav", voice=DEFAULT_VOICE):
    """
    Synthesize speech from text using the Kokoro-82M model.
    
    Args:
        text (str): The text to convert to speech
        output_path (str): Path to save the output audio file
        voice (str): The voice ID to use for synthesis
        
    Returns:
        str: Path to the generated audio file
    """
    # Generate speech using Kokoro pipeline
    generator = pipeline(text, voice=voice)
    
    # Kokoro returns a generator that yields tuples of (grapheme_slice, phoneme_slice, audio)
    # We'll concatenate all audio segments for the complete output
    audio_segments = []
    
    for _, (_, _, audio) in enumerate(generator):
        audio_segments.append(audio)
    
    # Concatenate all audio segments if there are multiple
    if len(audio_segments) > 1:
        final_audio = np.concatenate(audio_segments)
    else:
        final_audio = audio_segments[0] if audio_segments else np.array([])
    
    # Save the audio file (Kokoro uses 24000 Hz by default)
    sf.write(output_path, final_audio, 24000)
    return output_path