import os
import gradio as gr
import spaces
import torch

# 支持的语言选项
LANGUAGES = {
    "Auto Detect": "auto",
    "English": "English",
    "Chinese": "Chinese",
    "Traditional Chinese": "Traditional Chinese",
    "Russian": "Russian",
    "Japanese": "Japanese",
    "Korean": "Korean",
    "Spanish": "Spanish",
    "French": "French",
    "Portuguese": "Portuguese",
    "German": "German",
    "Italian": "Italian",
    "Thai": "Thai",
    "Vietnamese": "Vietnamese",
    "Indonesian": "Indonesian",
    "Malay": "Malay",
    "Arabic": "Arabic",
    "Hindi": "Hindi",
    "Hebrew": "Hebrew",
    "Burmese": "Burmese",
    "Tamil": "Tamil",
    "Urdu": "Urdu",
    "Bengali": "Bengali",
    "Polish": "Polish",
    "Dutch": "Dutch",
    "Romanian": "Romanian",
    "Turkish": "Turkish",
    "Khmer": "Khmer",
    "Lao": "Lao",
    "Cantonese": "Cantonese",
    "Czech": "Czech",
    "Greek": "Greek",
    "Swedish": "Swedish",
    "Hungarian": "Hungarian",
    "Danish": "Danish",
    "Finnish": "Finnish",
    "Ukrainian": "Ukrainian",
    "Bulgarian": "Bulgarian",
    "Serbian": "Serbian",
    "Telugu": "Telugu",
    "Afrikaans": "Afrikaans",
    "Armenian": "Armenian",
    "Assamese": "Assamese",
    "Asturian": "Asturian",
    "Basque": "Basque",
    "Belarusian": "Belarusian",
    "Bosnian": "Bosnian",
    "Catalan": "Catalan",
    "Cebuano": "Cebuano",
    "Croatian": "Croatian",
    "Egyptian Arabic": "Egyptian Arabic",
    "Estonian": "Estonian",
    "Galician": "Galician",
    "Georgian": "Georgian",
    "Gujarati": "Gujarati",
    "Icelandic": "Icelandic",
    "Javanese": "Javanese",
    "Kannada": "Kannada",
    "Kazakh": "Kazakh",
    "Latvian": "Latvian",
    "Lithuanian": "Lithuanian",
    "Luxembourgish": "Luxembourgish",
    "Macedonian": "Macedonian",
    "Maithili": "Maithili",
    "Maltese": "Maltese",
    "Marathi": "Marathi",
    "Mesopotamian Arabic": "Mesopotamian Arabic",
    "Moroccan Arabic": "Moroccan Arabic",
    "Najdi Arabic": "Najdi Arabic",
    "Nepali": "Nepali",
    "North Azerbaijani": "North Azerbaijani",
    "North Levantine Arabic": "North Levantine Arabic",
    "Northern Uzbek": "Northern Uzbek",
    "Norwegian Bokmål": "Norwegian Bokmål",
    "Norwegian Nynorsk": "Norwegian Nynorsk",
    "Occitan": "Occitan",
    "Odia": "Odia",
    "Pangasinan": "Pangasinan",
    "Sicilian": "Sicilian",
    "Sindhi": "Sindhi",
    "Sinhala": "Sinhala",
    "Slovak": "Slovak",
    "Slovenian": "Slovenian",
    "South Levantine Arabic": "South Levantine Arabic",
    "Swahili": "Swahili",
    "Tagalog": "Tagalog",
    "Ta’izzi-Adeni Arabic": "Ta’izzi-Adeni Arabic",
    "Tosk Albanian": "Tosk Albanian",
    "Tunisian Arabic": "Tunisian Arabic",
    "Venetian": "Venetian",
    "Waray": "Waray",
    "Welsh": "Welsh",
    "Western Persian": "Western Persian"
}

from huggingface_hub import snapshot_download
from vllm import LLM
from vllm.sampling_params import BeamSearchParams

save_dir = "./model_weights"
repo_id = "ByteDance-Seed/Seed-X-PPO-7B"
cache_dir = save_dir + "/cache"

snapshot_download(
  cache_dir=cache_dir,
  local_dir=save_dir,
  repo_id=repo_id,
  local_dir_use_symlinks=False,
  resume_download=True,
  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
)

device = torch.device("cuda")
model = LLM(model=save_dir,
                max_num_seqs=512,
                # tensor_parallel_size=8,
                enable_prefix_caching=True, 
                gpu_memory_utilization=0.95).to(device)

decoding_params = BeamSearchParams(beam_width=4,
                                   max_tokens=512)

@spaces.GPU(duration=120)
def translate_text(text, source_lang, target_lang):
    if not text.strip():
        return "请输入要翻译的文本"
    try:
        response = model.generate(messages, decoding_params)
        response = [res.outputs[0].text.strip() for res in results]
        yield response[0]
    except Exception as e:
        yield f"翻译出错: {str(e)}"


# 创建 Gradio 界面
with gr.Blocks(title="Qwen3-MT Translator") as demo:
    gr.Markdown("# 🌍 Seed-X")
    # gr.Markdown(
    #     'A real-time translation tool based on the Qwen3-MT model<br><a href="https://www.alibabacloud.com/help/en/model-studio/translation-abilities" target="_blank">Learn more about Qwen3-MT and API documentation</a>',
    #     elem_id="desc"
    # )
    gr.Image(
        value="https://github.com/ByteDance-Seed/Seed-X-7B/blob/main/imgs/LOGO.jpg",
        label="Seed-X",
        show_label=False,
        show_download_button=False,
        interactive=False,
        height=60
    )
    
    with gr.Row():
        with gr.Column():
            source_text = gr.Textbox(
                label="Input Text",
                placeholder="Please enter the text to translate...",
                lines=5
            )
            with gr.Row():
                source_lang = gr.Dropdown(
                    choices=list(LANGUAGES.keys()),
                    value="Auto Detect",
                    label="Source Language"
                )
                target_lang = gr.Dropdown(
                    choices=list(LANGUAGES.keys())[1:],  # Exclude "Auto Detect"
                    value="English",
                    label="Target Language"
                )
            translate_btn = gr.Button("Translate", variant="primary")
        
        with gr.Column():
            target_text = gr.Textbox(
                label="Translation Result",
                interactive=False,
                lines=5
            )
    
    # # 示例
    # gr.Examples(
    #     examples=[
    #         ["你好，世界！", "Chinese", "English"],
    #         ["Hello, how are you today?", "English", "Chinese"],
    #         ["私は学生です。", "Japanese", "Chinese"],
    #         ["Bonjour, comment allez-vous?", "French", "English"]
    #     ],
    #     inputs=[source_text, source_lang, target_lang],
    #     outputs=target_text,
    #     fn=translate_text,
    #     cache_examples=True
    # )
    
    # 按钮点击事件
    translate_btn.click(
        fn=translate_text,
        inputs=[source_text, source_lang, target_lang],
        outputs=target_text
    )
    
    # 支持回车键翻译
    source_text.submit(
        fn=translate_text,
        inputs=[source_text, source_lang, target_lang],
        outputs=target_text
    )

# 启动应用
if __name__ == "__main__":
    demo.launch()