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

# ================== 1. Emotion Presets ==================
EMOTION_PRESETS = {
    "😊 ပျော်ရွှင် (Cheerful)": {"pitch": 5, "rate": 7, "volume": "medium"},
    "😢 ဝမ်းနည်း (Sad)": {"pitch": -5, "rate": -10, "volume": "soft"},
    "😠 ဒေါသ (Angry)": {"pitch": 5, "rate": 12, "volume": "loud"},
    "😌 အေးဆေး (Calm)": {"pitch": -3, "rate": -5, "volume": "medium"},
    "🎉 စိတ်လှုပ်ရှား (Excited)": {"pitch": 7, "rate": 10, "volume": "loud"},
    "🔧 Custom (ကိုယ်တိုင်ပြင်ဆင်)": {"pitch": 0, "rate": 0, "volume": "medium"}
}

def update_emotion(emotion):
    preset = EMOTION_PRESETS[emotion]
    return gr.update(value=preset["pitch"]), gr.update(value=preset["rate"]), gr.update(value=preset["volume"])

# ================== 2. Audio Effects ==================
def apply_eq(audio, bass_gain_db, treble_gain_db):
    if bass_gain_db != 0:
        bass_part = audio.low_pass_filter(250).apply_gain(bass_gain_db)
        audio = audio.overlay(bass_part)
    if treble_gain_db != 0:
        treble_part = audio.high_pass_filter(4000).apply_gain(treble_gain_db)
        audio = audio.overlay(treble_part)
    return audio

def apply_reverb(audio, amount=50, delay_ms=80, echoes=5):
    if amount <= 0:
        return audio
    decay_db_per_echo = -amount / 15.0
    wet = AudioSegment.silent(duration=len(audio) + echoes * delay_ms)
    for i in range(1, echoes+1):
        gain = decay_db_per_echo * i
        echo = audio.apply_gain(gain)
        silence_before = AudioSegment.silent(duration=i*delay_ms)
        echo_with_delay = silence_before + echo
        wet = wet.overlay(echo_with_delay)
    combined = audio.overlay(wet[:len(audio)])
    return combined

# ================== 3. Main TTS Pipeline (No auto breaks) ==================
async def tts_pipeline(text, emotion, pitch, rate, volume_ssml,
                       bass_gain, treble_gain, reverb_amount, output_gain,
                       voice):
    pitch_str = f"+{pitch}%" if pitch >= 0 else f"{pitch}%"
    rate_str = f"+{rate}%" if rate >= 0 else f"{rate}%"

    ssml = f"""<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis">
    <voice name="{voice}">
        <prosody pitch="{pitch_str}" rate="{rate_str}" volume="{volume_ssml}">
            {text}
        </prosody>
    </voice>
</speak>"""

    communicate = edge_tts.Communicate(ssml, voice=voice)
    tmp_mp3 = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    await communicate.save(tmp_mp3.name)

    audio = AudioSegment.from_mp3(tmp_mp3.name)
    os.unlink(tmp_mp3.name)

    audio = apply_eq(audio, bass_gain, treble_gain)
    audio = apply_reverb(audio, reverb_amount)
    audio = audio.apply_gain(output_gain)

    out_mp3 = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    audio.export(out_mp3.name, format="mp3")
    return out_mp3.name

# ================== 4. Gradio Interface ==================
with gr.Blocks(title="Myanmar Emotional TTS") as demo:
    gr.Markdown("## 🎙️ မြန်မာ Emotional TTS (Edge TTS + Effects)")
    gr.Markdown("စာရိုက်ထည့်ပါ၊ Emotion ရွေးပါ၊ Pitch/Rate ကို ကိုယ်တိုင်ညှိပြီး အသံအကျိုးသက်ရောက်မှုများ (EQ, Reverb, Volume) ထည့်နိုင်ပါတယ်။")

    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(label="📝 မြန်မာစာသား", lines=5, placeholder="ဒီမှာ စာရိုက်ထည့်ပါ...")

        with gr.Column(scale=1):
            voice = gr.Dropdown(label="🎤 အသံ", choices=["my-MM-NilarNeural", "my-MM-ThihaNeural"], value="my-MM-NilarNeural")
            emotion = gr.Dropdown(label="🎭 Emotion", choices=list(EMOTION_PRESETS.keys()), value="😊 ပျော်ရွှင် (Cheerful)")

    with gr.Row():
        pitch = gr.Slider(label="🎵 Pitch (%)", minimum=-20, maximum=20, value=5, step=1)
        rate = gr.Slider(label="⏩ Rate (%)", minimum=-20, maximum=20, value=7, step=1)
        volume_ssml = gr.Dropdown(label="🔊 Volume (SSML)", choices=["x-soft","soft","medium","loud","x-loud"], value="medium")

    gr.Markdown("### 🎚️ Audio Effects")
    with gr.Row():
        bass_gain = gr.Slider(label="🎧 Bass Gain (dB)", minimum=-10, maximum=10, value=0, step=1)
        treble_gain = gr.Slider(label="🎧 Treble Gain (dB)", minimum=-10, maximum=10, value=0, step=1)
    with gr.Row():
        reverb_amount = gr.Slider(label="🌊 Reverb Amount", minimum=0, maximum=100, value=0, step=1)
        output_gain = gr.Slider(label="🔈 Output Volume (dB)", minimum=-10, maximum=20, value=0, step=1)

    generate_btn = gr.Button("▶️ အသံထုတ်မယ်", variant="primary")
    audio_output = gr.Audio(label="🔊 ရလာတဲ့ အသံ")

    emotion.change(
        fn=update_emotion,
        inputs=emotion,
        outputs=[pitch, rate, volume_ssml]
    )

    generate_btn.click(
        fn=tts_pipeline,
        inputs=[text_input, emotion, pitch, rate, volume_ssml,
                bass_gain, treble_gain, reverb_amount, output_gain,
                voice],
        outputs=audio_output
    )

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