#!/usr/bin/env python3
"""Gradio Space for MiMo-V2.5-ASR 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/mimo_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_audio.stt import load

MODEL_ID = os.environ.get("MIMO_MLX_MODEL", "mlx-community/MiMo-V2.5-ASR-MLX")
MODEL_DIR = os.environ.get("MIMO_MLX_MODEL_DIR", "/models/MiMo-V2.5-ASR-MLX")
AUDIO_TOKENIZER_DIR = os.environ.get(
    "MIMO_AUDIO_TOKENIZER_DIR",
    "/models/MiMo-Audio-Tokenizer",
)
LANGUAGES = {"Auto": None, "Chinese": "zh", "English": "en"}


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_asr_model():
    kwargs = {}
    if os.path.exists(os.path.join(AUDIO_TOKENIZER_DIR, "config.json")):
        kwargs["audio_tokenizer_dir"] = AUDIO_TOKENIZER_DIR
    return load(_model_source(), **kwargs)


def transcribe(audio_path: str | None, language_label: str, max_tokens: int):
    if not audio_path:
        return "", "Upload or record an audio file."

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

    infer_start = time.perf_counter()
    result = model.generate(
        audio_path,
        language=LANGUAGES.get(language_label),
        max_tokens=int(max_tokens),
    )
    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 result.text, status


with gr.Blocks(title="MiMo-V2.5-ASR MLX") as demo:
    gr.Markdown("# MiMo-V2.5-ASR MLX")
    gr.Markdown(
        "This Space uses MLX CPU fallback on Linux. It can verify that the model "
        "loads, but full ASR 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():
            audio = gr.Audio(label="Audio", type="filepath", sources=["upload", "microphone"])
            language = gr.Radio(
                choices=list(LANGUAGES.keys()),
                value="Auto",
                label="Language",
            )
            max_tokens = gr.Slider(1, 128, value=16, step=1, label="Max Tokens")
            button = gr.Button("Transcribe", variant="primary")
        with gr.Column():
            transcript = gr.Textbox(label="Transcript", lines=8)
            status = gr.Textbox(label="Status")

    button.click(transcribe, [audio, language, max_tokens], [transcript, status])


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