import spaces  # MUST come before any torch/CUDA-touching import
import torch
import gradio as gr

from transformers import AutoModelForImageTextToText, AutoTokenizer

MODEL_ID = "17slever17/translate-gemma-4-sub-e4b"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda")
model.eval()

SUPPORTED_LANGUAGES = [
    "English", "Russian", "Spanish", "German", "Japanese",
    "French", "Portuguese", "Chinese", "Dutch", "Italian", "Korean",
]

STYLES = ["neutral", "friendly", "official"]


def _build_system_message(source_lang: str, target_lang: str, rules: str, style: str) -> str:
    """Build the system message exactly as specified in the model card."""
    parts = [f"TASK: Translate {source_lang} subtitles into {target_lang}."]
    if rules.strip():
        parts.append(f"RULES: {rules.strip()}")
    parts.append(f"STYLE: {style}.")
    parts.append(
        "Translate only CURRENT_SOURCE. PREVIOUS_SOURCE and PREVIOUS_TRANSLATION are context only. "
        "Preserve meaning, tone, slang, profanity, uncertainty, repetitions and incomplete speech. "
        "Return only the final translation without labels or commentary."
    )
    return "\n".join(parts)


def _build_user_message(prev_source: str, prev_translation: str, current_source: str) -> str:
    """Build the user message following the model card's format."""
    blocks = []
    if prev_source.strip():
        blocks.append(f"[PREVIOUS_SOURCE]\n{prev_source.strip()}")
    if prev_translation.strip():
        blocks.append(f"[PREVIOUS_TRANSLATION]\n{prev_translation.strip()}")
    blocks.append(f"[CURRENT_SOURCE]\n{current_source.strip()}")
    return "\n\n".join(blocks)


@spaces.GPU(duration=60)
def translate(
    current_source: str,
    source_language: str = "English",
    target_language: str = "Russian",
    previous_source: str = "",
    previous_translation: str = "",
    rules: str = "",
    style: str = "friendly",
) -> str:
    """Translate a subtitle line using context-aware Gemma 4 Sub model.

    Args:
        current_source: The subtitle segment to translate.
        source_language: Language of the source text.
        target_language: Language to translate into.
        previous_source: Previous source-language subtitles for context (optional).
        previous_translation: Previous translated subtitles for context (optional).
        rules: Additional rules (speaker name, gender, terminology, etc.) (optional).
        style: Translation style — neutral, friendly, or official.
    """
    if not current_source.strip():
        return "Please enter a subtitle line to translate."

    system_msg = _build_system_message(source_language, target_language, rules, style)
    user_msg = _build_user_message(previous_source, previous_translation, current_source)

    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_msg},
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to("cuda")

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
        )

    generated_tokens = output[0, inputs["input_ids"].shape[1]:]
    translation = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
    return translation


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# 🌐 Translate Gemma 4 Sub — Context-Aware Subtitle Translation\n"
        "A multilingual Gemma 4 fine-tune specialized for subtitle translation. "
        "Provide previous source and translation context for natural, coherent translations."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            source_language = gr.Dropdown(
                label="Source Language",
                choices=SUPPORTED_LANGUAGES,
                value="English",
            )
            target_language = gr.Dropdown(
                label="Target Language",
                choices=SUPPORTED_LANGUAGES,
                value="Russian",
            )
            style = gr.Dropdown(
                label="Style",
                choices=STYLES,
                value="friendly",
            )

        current_source = gr.Textbox(
            label="Current Source (to translate)",
            placeholder="Yeah, well... I changed my mind.",
            lines=2,
        )

        with gr.Accordion("Context & Rules (optional)", open=False):
            previous_source = gr.Textbox(
                label="Previous Source (context only)",
                placeholder="I thought you said you weren't coming.",
                lines=2,
            )
            previous_translation = gr.Textbox(
                label="Previous Translation (context only)",
                placeholder="Я думала, ты сказала, что не придёшь.",
                lines=2,
            )
            rules = gr.Textbox(
                label="Rules (speaker name, gender, terminology, etc.)",
                placeholder="speaker gender: female",
                lines=2,
            )

        run_btn = gr.Button("Translate", variant="primary")
        output = gr.Textbox(
            label="Translation",
            lines=3,
            interactive=False,
        )

        run_btn.click(
            fn=translate,
            inputs=[
                current_source, source_language, target_language,
                previous_source, previous_translation, rules, style,
            ],
            outputs=output,
            api_name="translate",
        )

    gr.Examples(
        examples=[
            ["Yeah, well... I changed my mind.", "English", "Russian", "friendly"],
            ["I thought you said you weren't coming.", "English", "Spanish", "neutral"],
            ["Wait, hold on a second...", "English", "Japanese", "friendly"],
            ["That's not what I meant at all.", "English", "German", "official"],
            ["Could you repeat that, please?", "English", "French", "official"],
        ],
        inputs=[current_source, source_language, target_language, style],
        outputs=output,
        fn=translate,
        cache_examples=True,
        cache_mode="lazy",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)