import gradio as gr
import edge_tts
import asyncio
import tempfile

from piper_engine import PiperEngine


# -----------------------------
# Load Piper Once
# -----------------------------

piper_engine = PiperEngine(
    "models/en_US-amy-medium.onnx"
)


# -----------------------------
# Voice Loader (Edge)
# -----------------------------

async def get_voices():

    voices = await edge_tts.list_voices()

    return {
        f"{v['ShortName']} - {v['Locale']} ({v['Gender']})": v["ShortName"]
        for v in voices
    }


# -----------------------------
# Edge TTS (Neutral)
# -----------------------------

async def edge_tts_engine(text, voice, rate, pitch):

    if not text.strip():
        return None, "Please enter text."

    if not voice:
        return None, "Please select a voice."

    voice_short = voice.split(" - ")[0]

    rate_str = f"{rate:+d}%"
    pitch_str = f"{pitch:+d}Hz"

    communicate = edge_tts.Communicate(
        text,
        voice_short,
        rate=rate_str,
        pitch=pitch_str,
    )

    with tempfile.NamedTemporaryFile(
        delete=False,
        suffix=".mp3",
    ) as tmp:

        path = tmp.name
        await communicate.save(path)

    return path, None


# -----------------------------
# Piper TTS (Expressive)
# -----------------------------

def piper_engine_wrapper(text, rate, pitch):
    if not text.strip():
        return None, "Please enter text."

    speed = 1.0 + rate / 100      # rate from slider
    pitch_scale = 1.0 + pitch / 50
    audio = piper_engine.synthesize(text, speed, pitch_scale)

    return audio, None
    
# -----------------------------
# Unified Interface
# -----------------------------

async def unified_tts(
    text,
    voice,
    rate,
    pitch,
    engine,
):

    if engine == "Edge (Neutral)":

        return await edge_tts_engine(
            text,
            voice,
            rate,
            pitch,
        )

    else:

        return piper_engine_wrapper(
            text,
            rate,
            pitch,
        )


# -----------------------------
# UI Builder
# -----------------------------

async def create_demo():

    voices = await get_voices()

    with gr.Blocks(analytics_enabled=False) as demo:

        gr.Markdown("""
        # 🎙️ AI Text-to-Speech Lab
        Compare traditional and expressive AI voices.
        """)

        # Text
        text_input = gr.Textbox(
            label="Text",
            lines=6,
            placeholder="Paste lecture notes, narration, or scripts here...",
        )

        # Engine Selector
        engine_radio = gr.Radio(
            ["Edge (Neutral)", "Piper (Expressive)"],
            value="Edge (Neutral)",
            label="Generation",
        )

        # Voice (Edge Only)
        voice_dropdown = gr.Dropdown(
            choices=[""] + list(voices.keys()),
            label="Edge Voice",
            value="",
        )

        with gr.Row():

            rate_slider = gr.Slider(
                -50, 50, 0,
                step=1,
                label="Emotion / Speed",
            )

            pitch_slider = gr.Slider(
                -20, 20, 0,
                step=1,
                label="Intonation / Pitch",
            )

        generate_btn = gr.Button(
            "Generate Audio",
            variant="primary",
        )

        audio_output = gr.Audio(
            label="Output",
        )

        warning_md = gr.Markdown()


        generate_btn.click(
            fn=unified_tts,
            inputs=[
                text_input,
                voice_dropdown,
                rate_slider,
                pitch_slider,
                engine_radio,
            ],
            outputs=[
                audio_output,
                warning_md,
            ],
        )


    return demo


# -----------------------------
# Runner
# -----------------------------

async def main():

    demo = await create_demo()

    demo.queue(default_concurrency_limit=20)

    demo.launch()


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