import gradio as gr
import ollama

MODEL_NAME = "translategemma"

LANGUAGES = {
    "English": {"code": "en", "name": "English"},
    "Chinese (Simplified)": {"code": "zh-Hans", "name": "Chinese"},
    "Chinese (Traditional)": {"code": "zh-Hant", "name": "Chinese"},
    "Japanese": {"code": "ja", "name": "Japanese"},
    "Korean": {"code": "ko", "name": "Korean"},
    "French": {"code": "fr", "name": "French"},
    "Spanish": {"code": "es", "name": "Spanish"},
    "German": {"code": "de", "name": "German"},
    "Russian": {"code": "ru", "name": "Russian"},
}


def build_prompt(source_ui_name, target_ui_name, input_text):
    src_info = LANGUAGES[source_ui_name]
    tgt_info = LANGUAGES[target_ui_name]

    source_lang = src_info['name']
    source_code = src_info['code']
    target_lang = tgt_info['name']
    target_code = tgt_info['code']

    prompt = (
        f"You are a professional {source_lang} ({source_code}) to {target_lang} ({target_code}) translator. "
        f"Your goal is to accurately convey the meaning and nuances of the original {source_lang} text "
        f"while adhering to {target_lang} grammar, vocabulary, and cultural sensitivities.\n"
        f"Produce only the {target_lang} translation, without any additional explanations or commentary. "
        f"Please translate the following {source_lang} text into {target_lang}:\n"
        f"\n"
        f"\n"
        f"{input_text}"
    )
    
    return prompt


def translate(source_ui_name, target_ui_name, input_text):
    if not input_text or not input_text.strip():
        yield ""
        return

    full_prompt = build_prompt(source_ui_name, target_ui_name, input_text)

    try:
        stream = ollama.generate(
            model=MODEL_NAME,
            prompt=full_prompt,
            stream=True
        )

        partial_text = ""
        for chunk in stream:
            content = chunk.get("response", "")
            partial_text += content
            yield partial_text

    except Exception as e:
        yield f"Error: {str(e)}"


with gr.Blocks(title="TranslateGemma") as demo:
    gr.Markdown("# 🌍 TranslateGemma")

    with gr.Row():
        with gr.Column():
            source_dropdown = gr.Dropdown(
                choices=list(LANGUAGES.keys()),
                value="English",
                container=False
            )
            input_text = gr.Textbox(
                lines=10,
                placeholder="Enter text to translate here...",
                autofocus=True,
                container=False
            )

        with gr.Column():
            target_dropdown = gr.Dropdown(
                choices=list(LANGUAGES.keys()),
                value="Chinese (Simplified)",
                container=False
            )
            output_text = gr.Textbox(
                lines=10,
                container=False,
                interactive=False
            )

    translate_btn = gr.Button("Translate", variant="primary")

    translate_btn.click(
        fn=translate,
        inputs=[source_dropdown, target_dropdown, input_text],
        outputs=output_text
    )


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