import sys

import gradio as gr
import edge_tts
import asyncio
import tempfile
import os

from edge_tts import SubMaker


async def get_voices():
    voices = await edge_tts.list_voices()
    return {f"{v['ShortName']} - {v['Locale']} ({v['Gender']})": v['ShortName'] for v in voices}


async def text_to_speech(text, voice, rate, pitch):
    if not text.strip():
        return None, "Please enter text to convert."
    if not voice:
        return None, "Please select a voice."

    # 创建临时文件
    audio_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    audio_path = audio_tmp.name
    audio_tmp.close()  # 关闭文件句柄以便后续重新打开

    sub_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".srt")
    sub_path = sub_tmp.name
    sub_tmp.close()

    voice_short_name = voice.split(" - ")[0]
    rate_str = f"{rate:+d}%"
    pitch_str = f"{pitch:+d}Hz"
    communicate = edge_tts.Communicate(text, voice_short_name, rate=rate_str, pitch=pitch_str)
    submaker = SubMaker()

    audio_file = None
    sub_file = None
    try:
        # 打开临时文件进行写入
        audio_file = open(audio_path, "wb")
        sub_file = open(sub_path, "w", encoding="utf-8")

        async for chunk in communicate.stream():
            if chunk["type"] == "audio":
                audio_file.write(chunk["data"])
            elif chunk["type"] == "WordBoundary":
                submaker.feed(chunk)
        submaker.merge_cues(12)
        # 写入字幕内容
        sub_file.write(submaker.get_srt())

    except Exception as e:
        # 清理临时文件
        if os.path.exists(audio_path):
            os.remove(audio_path)
        if os.path.exists(sub_path):
            os.remove(sub_path)
        return None, None, str(e)
    finally:
        # 确保文件正确关闭
        if audio_file:
            audio_file.close()
        if sub_file:
            sub_file.close()

    return audio_path, sub_path, None

async def tts_interface(text, voice, rate, pitch):
    audio, srt, warning = await text_to_speech(text, voice, rate, pitch)
    if warning:
        return audio, srt, gr.Warning(warning)
    return audio, srt,  None

async def create_demo():
    voices = await get_voices()
    

    demo = gr.Interface(
        fn=tts_interface,
        inputs=[
            gr.Textbox(label="Input Text", lines=5),
            gr.Dropdown(choices=[""] + list(voices.keys()), label="选择配音员", value=""),
            gr.Slider(minimum=-50, maximum=50, value=0, label="Speech Rate Adjustment (%)", step=1),
            gr.Slider(minimum=-20, maximum=20, value=0, label="Pitch Adjustment (Hz)", step=1)
        ],
        outputs=[
            gr.Audio(label="Generated Audio", type="filepath"),
            gr.File(label="下载 SRT 字幕文件"),  # 直接提供 SRT 下载
            gr.Markdown(label="Warning", visible=False)
        ],
        title="Edge TTS Text-to-Speech",
        article="Experience the power of Edge TTS for text-to-speech conversion, and explore our advanced Text-to-Video Converter for even more creative possibilities!",
        analytics_enabled=False,
        allow_flagging="manual",
        api_name=None
    )
    return demo

async def main():
    demo = await create_demo()
    demo.queue(default_concurrency_limit=5)
    demo.launch(show_api=False)

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