#!/usr/bin/env python3
"""Gradio Space for Hy-MT2 MLX."""

from __future__ import annotations

import os
import time
from functools import lru_cache
from pathlib import Path

os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")


def _install_gxx_wrapper() -> None:
    """Make MLX CPU JIT tolerate GCC's built-in _Float* typedefs on Spaces."""
    wrapper_dir = Path("/tmp/hymt2_mlx_bin")
    wrapper_dir.mkdir(parents=True, exist_ok=True)
    wrapper = wrapper_dir / "g++"
    if not wrapper.exists():
        wrapper.write_text("#!/usr/bin/env bash\nexec /usr/bin/g++ -fpermissive \"$@\"\n", encoding="utf-8")
        wrapper.chmod(0o755)
    os.environ["PATH"] = f"{wrapper_dir}:{os.environ.get('PATH', '')}"


_install_gxx_wrapper()

import gradio as gr
from mlx_lm import generate, load
from mlx_lm.sample_utils import make_logits_processors, make_sampler

MODEL_ID = os.environ.get("HYMT2_MLX_MODEL", "mlx-community/Hy-MT2-1.8B-4bit")
MODEL_DIR = os.environ.get("HYMT2_MLX_MODEL_DIR", "/models/Hy-MT2-1.8B-4bit")
LANGUAGES = [
    "English", "Chinese", "Traditional Chinese", "Japanese", "Korean", "French",
    "German", "Spanish", "Portuguese", "Russian", "Arabic", "Hindi", "Vietnamese",
    "Thai", "Indonesian", "Turkish", "Ukrainian",
]


def _model_source() -> str:
    if os.path.exists(os.path.join(MODEL_DIR, "config.json")):
        return MODEL_DIR
    return MODEL_ID


@lru_cache(maxsize=1)
def load_translation_model():
    return load(_model_source())


def build_prompt(text: str, target_lang: str) -> str:
    return (
        f"Translate the following text into {target_lang}. "
        "Note that you should only output the translated result without any "
        "additional explanation:\n\n"
        f"{text.strip()}"
    )


def translate(text: str, target_lang: str, max_tokens: int):
    if not text or not text.strip():
        return "", "Enter text to translate."

    start = time.perf_counter()
    was_loaded = load_translation_model.cache_info().currsize > 0
    load_start = time.perf_counter()
    model, tokenizer = load_translation_model()
    load_elapsed = time.perf_counter() - load_start

    prompt = build_prompt(text, target_lang)
    if getattr(tokenizer, "has_chat_template", False):
        prompt = tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            tokenize=False,
            add_generation_prompt=True,
        )

    sampler = make_sampler(0.7, 0.6, 0.0, 1, 20)
    logits_processors = make_logits_processors(
        repetition_penalty=1.05,
        repetition_context_size=20,
    )
    infer_start = time.perf_counter()
    output = generate(
        model,
        tokenizer,
        prompt,
        max_tokens=int(max_tokens),
        sampler=sampler,
        logits_processors=logits_processors,
        verbose=False,
    )
    infer_elapsed = time.perf_counter() - infer_start
    total_elapsed = time.perf_counter() - start
    load_text = "cached" if was_loaded else f"{load_elapsed:.2f}s"
    status = (
        f"Done. elapsed={total_elapsed:.2f}s, "
        f"model_load={load_text}, infer={infer_elapsed:.2f}s"
    )
    return output.strip(), status


with gr.Blocks(title="Hy-MT2 MLX") as demo:
    gr.Markdown("# Hy-MT2 MLX")
    gr.Markdown(
        "This Space uses MLX CPU fallback on Linux. It can verify that the model "
        "loads, but full translation on cpu-basic may exceed the request timeout. "
        "Run locally on Apple Silicon for the intended MLX GPU path."
    )
    with gr.Row():
        with gr.Column():
            source = gr.Textbox(
                label="Source Text",
                lines=6,
                value="今天天气真好。",
            )
            target = gr.Dropdown(LANGUAGES, value="English", label="Target Language")
            max_tokens = gr.Slider(1, 128, value=16, step=1, label="Max Tokens")
            button = gr.Button("Translate", variant="primary")
        with gr.Column():
            output = gr.Textbox(label="Translation", lines=8)
            status = gr.Textbox(label="Status")

    button.click(translate, [source, target, max_tokens], [output, status])
    gr.Examples(
        examples=[
            ["今天天气真好。", "English", 64],
            ["The meeting has been moved to next Monday.", "Chinese", 64],
        ],
        inputs=[source, target, max_tokens],
    )


if __name__ == "__main__":
    demo.queue(max_size=4, default_concurrency_limit=1).launch()