import tempfile
import edge_tts
import gradio as gr
import asyncio

language_dict = {
    "Amharic": {
        "Ameha": "am-ET-AmehaNeural",
        "Mekdes": "am-ET-MekdesNeural"
    },
    "English": {
        "Ryan": "en-GB-RyanNeural",
        "Clara": "en-CA-ClaraNeural"
    }
}

async def text_to_speech_edge(text, language, speaker):
    # Handle Tigrinya as Amharic with different label
    if language == "Tigrinya":
        language = "Amharic"
    
    try:
        voice = language_dict[language][speaker]
    except KeyError:
        error_msg = (f"ስህተት: ድምፅ '{speaker}' አልተገኘም።"
                    if language == "Amharic" 
                    else f"Error: Voice '{speaker}' not found")
        raise gr.Error(error_msg)

    try:
        communicate = edge_tts.Communicate(text, voice)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
            tmp_path = tmp_file.name
            await asyncio.wait_for(communicate.save(tmp_path), timeout=30)
        return tmp_path
        
    except asyncio.TimeoutError:
        error_msg = ("ስህተት: ጊዜ አልቋል። እባክዎ እንደገና ይሞክሩ።"
                    if language == "Amharic"
                    else "Error: Timeout. Please try again.")
        raise gr.Error(error_msg)
    except Exception as e:
        error_msg = (f"ስህተት: {str(e)}"
                    if language == "Amharic"
                    else f"Error: {str(e)}")
        raise gr.Error(error_msg)

def update_speakers(language):
    if language == "Tigrinya":
        speakers = list(language_dict["Amharic"].keys())
    else:
        speakers = list(language_dict.get(language, []))
    
    return gr.Dropdown(choices=speakers, value=speakers[0] if speakers else None)

with gr.Blocks(title="SelamGPT TTS") as demo:
    gr.Markdown("## Amharic, English & Tigrinya Text-to-Speech")
    
    with gr.Row():
        language = gr.Dropdown(
            choices=["Amharic", "English", "Tigrinya"],
            value="Amharic",
            label="Language"
        )
        speaker = gr.Dropdown(
            label="Speaker",
            allow_custom_value=False
        )
        input_text = gr.Textbox(label="Input Text")
    
    generate_btn = gr.Button("Generate Audio")
    output_audio = gr.Audio(label="Output")
    
    # API endpoints
    demo.load(
        fn=lambda: gr.Dropdown(choices=language_dict["Amharic"].keys()),
        outputs=speaker
    )
    
    language.change(
        update_speakers,
        inputs=language,
        outputs=speaker
    )
    
    generate_btn.click(
        text_to_speech_edge,
        inputs=[input_text, language, speaker],
        outputs=output_audio
    )

if __name__ == "__main__":
    demo.launch(server_port=7860)