import os

# ZeroGPU allocconf: chunked transcriptions allocate variable-size feature
# batches; keep transient spikes from fragmenting the caching allocator.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import shutil
import tempfile
import time
from pathlib import Path

import spaces  # MUST come before torch / transformers
import torch
import gradio as gr
from transformers import AutoModel, AutoProcessor

MODEL_ID = "oruk/oruk-loop-153M"

print("Loading processor ...", flush=True)
processor = AutoProcessor.from_pretrained(MODEL_ID)

print("Loading oruk-loop-153M ...", flush=True)
t0 = time.perf_counter()
# Authors' release path: trust_remote_code runs modeling_oruk_loop.py, which
# defaults to bf16 + merge=True on GPU (per-depth loop weights materialised
# for fast decoding, ~4.2 GB of VRAM).
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, device="cuda")
print(f"Model ready in {time.perf_counter() - t0:.1f}s", flush=True)


def _normalize_audio(audio_path: str, scratch_dir: str) -> str:
    """Convert browser-supplied audio (webm/mp4/mp3/stereo) to 16 kHz mono WAV.

    Microphone recordings arrive as .webm, which soundfile cannot read;
    pydub (ffmpeg) handles every container a browser can produce. Files that
    are already 16 kHz mono wav are passed through untouched, so bundled
    examples keep their exact bytes.
    """
    from pydub import AudioSegment

    audio = AudioSegment.from_file(audio_path)
    changed = False
    if audio.frame_rate != 16000:
        audio = audio.set_frame_rate(16000)
        changed = True
    if audio.channels > 1:
        audio = audio.set_channels(1)
        changed = True
    if not changed and Path(audio_path).suffix.lower() == ".wav":
        return audio_path
    out = Path(scratch_dir) / (Path(audio_path).stem + "_16k_mono.wav")
    audio.export(out, format="wav")
    return out.as_posix()


def _audio_seconds(audio_path) -> float:
    try:
        from pydub import AudioSegment

        return AudioSegment.from_file(audio_path).duration_seconds
    except Exception:
        return 0.0


def _estimate_duration(audio_path, *args, **kwargs) -> int:
    """Callable ZeroGPU duration, sized from measured runs on this Space.

    Measured: 0.84 s for 11 s of audio (beam 4) and 4.08 s for 38 s
    (beam 4, two chunks) — about 0.12 s per second of audio, plus ~2 s of
    weight streaming on the first request after boot. The budget below is
    that with margin (0.35 s/s + 15 s floor), capped for very long uploads.
    """
    if not audio_path:
        return 30
    seconds = _audio_seconds(audio_path)
    return min(240, max(30, 15 + int(seconds * 0.35)))


@spaces.GPU(duration=_estimate_duration)
def transcribe(audio_path: str, num_beams: int = 1) -> str:
    """Transcribe English speech to text with oruk-loop-153M.

    A 153M-parameter looped-transformer ASR model distilled from
    Qwen3-ASR-1.7B: one encoder block applied 24 times and one decoder
    block applied 28 times. Recordings longer than 30 s are automatically
    cut at quiet points into chunks and the chunk transcripts are joined.

    Args:
        audio_path: Path to an audio file (upload or microphone recording).
        num_beams: Beam size for decoding. 1 = greedy, 4 = the setting of
            the model card's headline WER numbers.

    Returns:
        The transcribed text, with casing and punctuation.
    """
    if not audio_path:
        raise gr.Error("Provide an audio file or record from your microphone first.")

    scratch = tempfile.mkdtemp(prefix="oruk_loop_")
    try:
        wav_path = _normalize_audio(audio_path, scratch)
        beams = max(1, int(num_beams))
        start = time.perf_counter()
        text = model.transcribe(wav_path, processor, num_beams=beams)
        elapsed = time.perf_counter() - start
        seconds = _audio_seconds(wav_path)
        print(
            f"TRANSCRIBE {elapsed:.2f}s for {seconds:.1f}s audio (beam {beams}) "
            f"-> {text[:120]!r}",
            flush=True,
        )
        return text
    finally:
        shutil.rmtree(scratch, ignore_errors=True)


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

EXAMPLE_FILES = [
    "jfk.wav",
    "librispeech_sample.wav",
    "eu_toxicity_25s.wav",
    "eu_ireland_38s.wav",
]

with gr.Blocks(title="oruk-loop-153M") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎙️ oruk-loop-153M — English ASR with a looped transformer\n\n"
            "**oruk-loop-153M** ([oruk/oruk-loop-153M](https://huggingface.co/oruk/oruk-loop-153M)) "
            "is an English speech recognizer with **153M stored parameters**, distilled from "
            "[Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B-hf) (13.3× larger). "
            "Each transformer stack is a single shared block run in a loop — one encoder block "
            "applied 24× and one decoder block applied 28× — with low-rank per-depth adapters. "
            "Mean WER on the Open ASR Leaderboard's seven English test sets: **7.12%** with beam 4.\n\n"
            "Upload or record audio (recordings longer than 30 s are chunked automatically) "
            "and get a transcript with casing and punctuation."
        )

        with gr.Row():
            with gr.Column():
                audio_in = gr.Audio(
                    label="Audio",
                    type="filepath",
                    sources=["upload", "microphone"],
                )
                transcribe_btn = gr.Button("Transcribe", variant="primary")
            with gr.Column():
                text_out = gr.Textbox(
                    label="Transcript",
                    lines=10,
                    placeholder="The recognized text will appear here ...",
                )

        with gr.Accordion("Decoding settings", open=False):
            beams = gr.Radio(
                choices=[("Greedy (fast)", 1), ("Beam 4 (best accuracy)", 4)],
                value=1,
                type="value",
                label="Decoding",
                info="The model card's headline numbers use beam 4; greedy is faster.",
            )

        gr.Examples(
            examples=[[f] for f in EXAMPLE_FILES],
            inputs=audio_in,
            fn=transcribe,
            outputs=text_out,
            label="Or try one of these",
            cache_examples=True,
            cache_mode="lazy",
        )

    transcribe_btn.click(
        transcribe,
        inputs=[audio_in, beams],
        outputs=text_out,
        api_name="transcribe",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)