import gradio as gr
import edge_tts
import asyncio
import os
import uuid
from pydub import AudioSegment

# Fetch available voices
async def get_voices():
    voices = await edge_tts.VoicesManager.create()
    voice_list = voices.find(Language="en")
    return {f"{v['FriendlyName']}": v['ShortName'] for v in voice_list}

# The core TTS function with Progression
async def text_to_speech_progress(text, voice_name, voice_dict, progress=gr.Progress()):
    if not text.strip():
        return None
    
    short_name = voice_dict[voice_name]
    
    # Split text into sentences or chunks for the progress bar to have "steps"
    chunks = [c.strip() for c in text.split('.') if c.strip()]
    if not chunks: chunks = [text] # Fallback for single short sentences
    
    combined_audio = AudioSegment.empty()
    session_id = str(uuid.uuid4())
    os.makedirs(session_id, exist_ok=True)
    
    progress(0, desc="Starting AI Voice Synthesis...")
    
    for i, chunk in enumerate(chunks):
        # Update the progress bar
        percent = (i + 1) / len(chunks)
        progress(percent, desc=f"Processing sentence {i+1} of {len(chunks)}...")
        
        chunk_path = os.path.join(session_id, f"chunk_{i}.mp3")
        communicate = edge_tts.Communicate(chunk, short_name)
        await communicate.save(chunk_path)
        
        # Merge audio
        segment = AudioSegment.from_mp3(chunk_path)
        combined_audio += segment
        
        # Clean up the small chunk immediately to save space
        os.remove(chunk_path)

    final_path = f"tts_{session_id}.mp3"
    combined_audio.export(final_path, format="mp3")
    os.rmdir(session_id)
    
    return final_path

# Gradio Setup
async def main():
    voice_dict = await get_voices()
    voice_choices = list(voice_dict.keys())

    def process(text, voice, progress=gr.Progress()):
        # Running the async function inside the synchronous Gradio wrapper
        return asyncio.run(text_to_speech_progress(text, voice, voice_dict, progress))

    with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as demo:
        gr.Markdown("# 🎙️ AI Voice Generator with Progress Tracker")
        gr.Markdown("Type your text below. The progress bar will track each sentence being processed.")
        
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(
                    label="Input Text", 
                    placeholder="Enter long text here to see the progress bar in action...", 
                    lines=8
                )
                voice_dropdown = gr.Dropdown(
                    choices=voice_choices, 
                    label="Select AI Voice", 
                    value=voice_choices[0]
                )
                submit_btn = gr.Button("Generate Voice", variant="primary")
            
            with gr.Column():
                audio_output = gr.Audio(label="Resulting Audio")
        
        # Connect the button to the function and pass the progress object
        submit_btn.click(
            fn=process, 
            inputs=[input_text, voice_dropdown], 
            outputs=audio_output
        )

    demo.launch()

if __name__ == "__main__":
    asyncio.run(main())