import gradio as gr
import srt
from datetime import timedelta
import openai
import os

# 配置OpenAI API（用户需要自行添加API密钥）
openai.api_key = os.getenv('OPENAI_API_KEY')

def translate_text(text, target_language):
    client = openai.OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"你是一个专业翻译，请准确将字幕内容翻译成{target_language}，保持口语化表达。"},
            {"role": "user", "content": text}
        ]
    )
    return response.choices[0].message.content

def process_srt(input_file, target_language):
    with open(input_file.name, 'r', encoding='utf-8') as f:
        subs = list(srt.parse(f.read()))
    
    translated_subs = []
    for sub in subs:
        translated_text = translate_text(sub.content, target_language)
        translated_subs.append(srt.Subtitle(
            index=sub.index,
            start=sub.start,
            end=sub.end,
            content=translated_text
        ))
    
    return srt.compose(translated_subs)

# Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("## SRT translation")
    with gr.Row():
        input_file = gr.File(label="Upload SRT file", type="filepath")
        target_language = gr.Dropdown(
            choices=["English", "Finnish", "Chinese", "Japanese", "Korean", "German", "French", "Spanish", "Italian", "Arabic", "Russian"],
            label="Chose target language",
            value="Finnish"
        )
        btn = gr.Button("Start Translation", variant="primary", scale=1)

    with gr.Row():
        preview_text = gr.Textbox(label="Preview", interactive=False)
        output_file = gr.File(label="Download translated file", visible=True)

    def process_file(file, target_language):
        translated = process_srt(file, target_language)
        output_path = os.path.join(os.getcwd(), "translated.srt")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(translated)
        return translated, output_path
    
    btn.click(
        fn=process_file,
        inputs=[input_file, target_language],
        outputs=[preview_text, output_file]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv('PORT', 7860)))