import gradio as gr
import asyncio
import base64
import os
import struct
from pydub import AudioSegment

MAX_CHARS = 50000

# ============================================================
# ЭДГЭЭР ФУНКЦУУД ТАНАЙ ХУУЧИН app.py-АС ЯГ ТАГ ХУУЛАГДСАН
# (логик ӨӨРЧЛӖГДӖГҺЙ, зөвхөн st.error -> raise gr.Error болсон)
# ============================================================

def change_audio_properties(file_path, rate, pitch, volume):
    audio = AudioSegment.from_file(file_path)
    if volume != 0:
        audio = audio + volume
    if rate != 0:
        speed_factor = 1 + (rate / 100.0)
        audio = audio._spawn(audio.raw_data, overrides={"frame_rate": int(audio.frame_rate * speed_factor)}).set_frame_rate(audio.frame_rate)
    if pitch != 0:
        pitch_factor = 1 + (pitch / 100.0)
        audio = audio._spawn(audio.raw_data, overrides={"frame_rate": int(audio.frame_rate * pitch_factor)}).set_frame_rate(audio.frame_rate)
    out_path = "processed_audio.mp3"
    audio.export(out_path, format="mp3")
    return out_path

async def run_edge_tts_async(text, voice_name, rate, pitch, volume):
    from edge_tts import Communicate
    v = "mn-MN-BataaNeural" if voice_name == "Батаа (эрэгтэй)" else "mn-MN-YesuiNeural"
    c = Communicate(text, v, rate=f"{rate:+d}%", pitch=f"{pitch:+d}Hz", volume=f"{volume:+d}%")
    await c.save("out_audio.mp3")

def run_edge_tts(text, voice_name, rate, pitch, volume):
    asyncio.run(run_edge_tts_async(text, voice_name, rate, pitch, volume))
    return "out_audio.mp3"

VOICE_MN_TO_API = {
    "Лхагваа (эрэгтэй)": "Charon", "Доржоо (эрэгтэй)": "Fenrir", "Батбаяр (эрэгтэй)": "Puck",
    "Мөнхбат (эрэгтэй)": "Orbit", "Энхбаяр (эрэгтэй)": "Zephyr", "Ганбаатар (эрэгтэй)": "Schedar",
    "Дулмаа (эмэгтэй)": "Aoede", "Номин (эмэгтэй)": "Kale", "Нарантуяа (эмэгтэй)": "Leda"
}

def run_gemini_tts(text, voice_name_mn, rate, pitch, volume):
    import requests
    api_key = os.environ.get("GEMINI_API_KEY", "")
    if not api_key:
        raise gr.Error("⚠️ GEMINI_API_KEY тохируулаагүй байна!")
    voice_name = VOICE_MN_TO_API.get(voice_name_mn, "Charon")
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-preview-tts:generateContent?key={api_key}"
    data = {
        "contents": [{"parts": [{"text": text}]}],
        "generationConfig": {
            "responseModalities": ["AUDIO"],
            "speechConfig": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": voice_name}}}
        }
    }
    response = requests.post(url, json=data, timeout=60)
    if response.status_code != 200:
        raise gr.Error(f"Gemini API Алдаа: {response.status_code}")
    result = response.json()
    audio_b64 = result["candidates"][0]["content"]["parts"][0]["inlineData"]["data"]
    audio_bytes = base64.b64decode(audio_b64)
    raw_wav = "out_audio.wav"
    with open(raw_wav, "wb") as f:
        channels, sample_rate, bits = 1, 24000, 16
        data_size = len(audio_bytes)
        header = struct.pack('<4sI4s4sIHHIIHH4sI', b'RIFF', data_size + 36, b'WAVE', b'fmt ', 16, 1, channels, sample_rate, sample_rate * channels * bits // 8, channels * bits // 8, bits, b'data', data_size)
        f.write(header + audio_bytes)
    return change_audio_properties(raw_wav, rate, pitch, volume)

# ============================================================
# ЭНЭ ФУНКЦ - Cloudflare-ээс /api/tts дуудах гол ЦЭГ.
# Gradio автоматаар "/gradio_api/call/generate_audio" эндпоинт
# үүсгэнэ, учир нь функцийн нэр "generate_audio".
# ============================================================
def generate_audio(text, engine_select, voice_select, voice_gemini, rate, pitch, volume):
    if not text or not text.strip():
        raise gr.Error("Текст хоосон байна")
    if len(text) > MAX_CHARS:
        raise gr.Error(f"Текст {MAX_CHARS} тэмдэгтээс ихгуй байх ёстой")

    if engine_select == "Edge TTS (Батаа / Есүй)":
        audio_path = run_edge_tts(text, voice_select, int(rate), int(pitch), int(volume))
    else:
        audio_path = run_gemini_tts(text, voice_gemini, int(rate), int(pitch), int(volume))

    return audio_path


# ============================================================
# GRADIO UI — Streamlit-ийн оронд. Дотоод логик ӨӨРЧЛӖГДӖГҺЙ.
# ============================================================
with gr.Blocks(title="КиноЭзэн - Дуу Хоолой Үүсгэгч") as demo:
    gr.Markdown("# 🎙️ Монгол Дуу Хоолой & Орчуулга")

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="📝 Эх Текст оруулах",
                placeholder="Текстээ оруулна уу...",
                lines=8,
                max_lines=20
            )
            engine_select = gr.Radio(
                ["Gemini TTS (Байгалийн)", "Edge TTS (Батаа / Есүй)"],
                label="🤖 Дууны Хөдөлгүүр",
                value="Edge TTS (Батаа / Есүй)"
            )

        with gr.Column():
            voice_select = gr.Dropdown(
                ["Батаа (эрэгтэй)", "Есүй (эмэгтэй)"],
                label="🎙️ Edge Хоолой сонгох",
                value="Батаа (эрэгтэй)",
                visible=True
            )
            voice_gemini = gr.Dropdown(
                list(VOICE_MN_TO_API.keys()),
                label="🎙️ Gemini Хоолой сонгох",
                value="Лхагваа (эрэгтэй)",
                visible=False
            )
            rate_slider = gr.Slider(-50, 50, value=0, step=5, label="⚡ Хурд (%)")
            pitch_slider = gr.Slider(-50, 50, value=0, step=1, label="🎵 Бүдүүн / Нарийн (Hz)")
            volume_slider = gr.Slider(-20, 20, value=0, step=1, label="🔊 Дууны чанга сул (dB)")
            generate_btn = gr.Button("🎬 ДУУ ХОЛОЙ ҮҮСГЭХ", variant="primary")

    audio_output = gr.Audio(label="Үр дүн")

    # Engine солих үед зөв dropdown-ийг харуулах/нуух
    def toggle_voice_visibility(engine):
        if engine == "Edge TTS (Батаа / Есүй)":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    engine_select.change(
        fn=toggle_voice_visibility,
        inputs=[engine_select],
        outputs=[voice_select, voice_gemini]
    )

    generate_btn.click(
        fn=generate_audio,
        inputs=[text_input, engine_select, voice_select, voice_gemini, rate_slider, pitch_slider, volume_slider],
        outputs=[audio_output],
        api_name="generate_audio"  # Энэ нэрээр /gradio_api/call/generate_audio endpoint үүснэ
    )

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