"""VibeVoice-ASR-BitNet — multilingual speech recognition running fully on CPU.

The model (microsoft/VibeVoice-ASR-BitNet) is a heterogeneously quantized ASR
system: an I8_S acoustic/semantic VAE tokenizer plus a ternary (I2_S) BitNet
Qwen2.5-1.5B decoder. It has no PyTorch/CUDA inference path — the only runtime
is the authors' ggml-based VibeASR.cpp engine, which this Space compiles from
source at startup and then drives through its `asr_infer` CLI.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
from huggingface_hub import hf_hub_download

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MODEL_REPO = "microsoft/VibeVoice-ASR-BitNet"
VAE_FILE = "vibeasr-vae-encoder-i8_s.gguf"
LM_FILE = "vibeasr-lm-i2_s-embed-q6_k.gguf"
ENGINE_GIT_URL = "https://github.com/microsoft/VibeASR.cpp.git"

WORK_DIR = Path(os.environ.get("VIBEASR_WORK", "/tmp/vibeasr"))
SRC_DIR = WORK_DIR / "VibeASR.cpp"
ASR_BIN = SRC_DIR / "build" / "bin" / "asr_infer"

TARGET_SR = 24000  # the tokenizer operates at 24 kHz
MAX_AUDIO_SECONDS = 40.0  # keep CPU latency reasonable
MAX_TOKENS = 2048
CTX_SIZE = 16384
BATCH_SIZE = 2048

# Streaming is done by transcribing the audio in sequential windows and
# emitting each window's text as soon as it is ready, so the transcript fills
# in incrementally instead of appearing all at once at the end.
STREAM_CHUNK_SECONDS = 8.0

# Dictation mode accumulates live microphone audio and re-transcribes once a
# little new speech has arrived, showing a running transcription of what the
# user is saying.
DICTATION_STEP_SECONDS = 3.0
DICTATION_MAX_SECONDS = 60.0


def _n_cpu() -> int:
    """Return the number of usable CPU cores, honouring cgroup quotas."""
    quota = None
    for path in ("/sys/fs/cgroup/cpu.max",):
        try:
            with open(path) as fh:
                parts = fh.read().split()
            if len(parts) == 2 and parts[0] != "max":
                quota = max(1, int(float(parts[0]) / float(parts[1])))
        except OSError:
            pass
    try:
        available = len(os.sched_getaffinity(0))
    except (AttributeError, OSError):
        available = os.cpu_count() or 2
    return max(1, min(quota or available, available))


N_CPU = _n_cpu()
MAX_THREADS = max(2, min(8, N_CPU * 2))


# ---------------------------------------------------------------------------
# One-time setup: compile the ggml engine, fetch the GGUF weights
# ---------------------------------------------------------------------------

def _run(cmd: list[str], cwd: Path | None = None, label: str = "") -> None:
    """Run a subprocess, echoing a trimmed log tail if it fails."""
    print(f"[setup] {label or cmd[0]}: {' '.join(str(c) for c in cmd)}", flush=True)
    proc = subprocess.run(
        [str(c) for c in cmd], cwd=str(cwd) if cwd else None,
        text=True, capture_output=True,
    )
    if proc.returncode != 0:
        print(proc.stdout[-4000:], flush=True)
        print(proc.stderr[-4000:], flush=True)
        raise RuntimeError(f"{label or cmd[0]} failed with exit code {proc.returncode}")


def build_engine() -> None:
    """Clone and compile VibeASR.cpp (`asr_infer`) into WORK_DIR."""
    if ASR_BIN.exists():
        print("[setup] engine already built", flush=True)
        return

    started = time.time()
    WORK_DIR.mkdir(parents=True, exist_ok=True)

    if not (SRC_DIR / "CMakeLists.txt").exists():
        shutil.rmtree(SRC_DIR, ignore_errors=True)
        _run(["git", "clone", "--depth", "1", ENGINE_GIT_URL, SRC_DIR], label="git clone")
        _run(
            ["git", "submodule", "update", "--init", "--recursive", "--depth", "1"],
            cwd=SRC_DIR, label="git submodule",
        )

    cmake = shutil.which("cmake") or "cmake"
    _run(
        [
            cmake, "-B", "build",
            "-DCMAKE_BUILD_TYPE=Release",
            "-DLLAMA_BUILD_TESTS=OFF",
            "-DLLAMA_BUILD_EXAMPLES=OFF",
            "-DLLAMA_BUILD_SERVER=OFF",
        ],
        cwd=SRC_DIR, label="cmake configure",
    )
    _run(
        [cmake, "--build", "build", "--target", "asr_infer", "-j", str(N_CPU)],
        cwd=SRC_DIR, label="cmake build",
    )

    if not ASR_BIN.exists():
        raise RuntimeError(f"build finished but {ASR_BIN} is missing")
    print(f"[setup] engine built in {time.time() - started:.1f}s", flush=True)


def fetch_weights() -> tuple[str, str]:
    """Download the two GGUF checkpoints and return their local paths."""
    started = time.time()
    local = os.environ.get("VIBEASR_MODEL_DIR")
    if local and (Path(local) / VAE_FILE).exists():
        return str(Path(local) / VAE_FILE), str(Path(local) / LM_FILE)
    vae = hf_hub_download(MODEL_REPO, VAE_FILE)
    lm = hf_hub_download(MODEL_REPO, LM_FILE)
    print(f"[setup] weights ready in {time.time() - started:.1f}s", flush=True)
    return vae, lm


print(f"[setup] detected {N_CPU} usable CPU core(s)", flush=True)
build_engine()
VAE_PATH, LM_PATH = fetch_weights()


# ---------------------------------------------------------------------------
# Audio preprocessing
# ---------------------------------------------------------------------------

def _load_mono_24k(audio_path: str) -> np.ndarray:
    """Decode any libsndfile-readable file to mono float32 at 24 kHz."""
    data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
    data = data.mean(axis=1)
    if data.size == 0:
        raise gr.Error("The uploaded audio appears to be empty.")
    return _resample_24k(data, sr)


def _resample_24k(data: np.ndarray, sr: int) -> np.ndarray:
    """Linearly resample mono float32 audio to 24 kHz."""
    if sr == TARGET_SR:
        return data.astype(np.float32)
    n_out = max(1, int(round(len(data) * TARGET_SR / sr)))
    return np.interp(
        np.linspace(0.0, len(data) - 1, n_out),
        np.arange(len(data), dtype=np.float64),
        data.astype(np.float64),
    ).astype(np.float32)


def _write_wav(data: np.ndarray) -> str:
    """Write mono float32 audio at 24 kHz to a temporary PCM16 WAV file."""
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    tmp.close()
    sf.write(tmp.name, data, TARGET_SR, subtype="PCM_16")
    return tmp.name


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def _asr_engine(
    data: np.ndarray,
    hotwords: str = "",
    n_threads: int = N_CPU,
    greedy: bool = True,
    temperature: float = 0.7,
    top_p: float = 0.9,
) -> str:
    """Run the ggml ASR engine on a mono 24 kHz array and return the transcript."""
    if data.size == 0:
        return ""
    wav_path = _write_wav(data)
    cmd = [
        str(ASR_BIN),
        "--vae-model", VAE_PATH,
        "--lm-model", LM_PATH,
        "--audio", wav_path,
        "-t", str(int(n_threads)),
        "-c", str(CTX_SIZE),
        "-b", str(BATCH_SIZE),
        "--max-tokens", str(MAX_TOKENS),
        "--prompt-format", "text",
    ]
    if greedy:
        cmd.append("--greedy")
    else:
        cmd += ["--temperature", str(float(temperature)), "--top-p", str(float(top_p))]
    if hotwords and hotwords.strip():
        cmd += ["--context", hotwords.strip()]

    try:
        proc = subprocess.run(cmd, text=True, capture_output=True, timeout=900)
    except subprocess.TimeoutExpired:
        raise gr.Error("Inference timed out. Please try a shorter clip.")
    finally:
        try:
            os.unlink(wav_path)
        except OSError:
            pass

    stdout = (proc.stdout or "").strip()
    stderr = (proc.stderr or "").strip()
    print(stderr[-3000:], flush=True)
    if proc.returncode != 0:
        raise gr.Error(f"Engine failed (exit {proc.returncode}): {stderr[-400:]}")
    return stdout


def transcribe(
    audio_path: str | None,
    hotwords: str = "",
    n_threads: int = N_CPU,
    greedy: bool = True,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """Transcribe speech with VibeVoice-ASR-BitNet on CPU, streaming as it goes.

    The audio is processed in sequential windows so the transcript is emitted
    incrementally — each window's text appears as soon as it is decoded, rather
    than waiting for the whole clip to finish.

    Args:
        audio_path: Path to an audio file (any format libsndfile can read).
        hotwords: Optional comma-separated context terms to bias decoding.
        n_threads: Number of CPU threads handed to the ggml engine.
        greedy: Use deterministic greedy decoding instead of sampling.
        temperature: Sampling temperature, only used when `greedy` is False.
        top_p: Nucleus sampling threshold, only used when `greedy` is False.

    Yields:
        Tuple of (transcript so far, markdown status), updated per window.
    """
    if not audio_path:
        raise gr.Error("Please upload or record some audio first.")

    data = _load_mono_24k(audio_path)
    max_samples = int(MAX_AUDIO_SECONDS * TARGET_SR)
    trimmed = len(data) > max_samples
    if trimmed:
        data = data[:max_samples]
    duration = len(data) / TARGET_SR

    chunk = int(STREAM_CHUNK_SECONDS * TARGET_SR)
    n_chunks = max(1, int(np.ceil(len(data) / chunk)))
    started = time.time()
    pieces: list[str] = []

    for i in range(n_chunks):
        segment = data[i * chunk:(i + 1) * chunk]
        yield (
            " ".join(pieces).strip(),
            f"Transcribing on CPU… window {i + 1}/{n_chunks}",
        )
        text = _asr_engine(segment, hotwords, n_threads, greedy, temperature, top_p)
        if text:
            pieces.append(text)
        yield (
            " ".join(pieces).strip(),
            f"Transcribing on CPU… window {i + 1}/{n_chunks}",
        )

    elapsed = time.time() - started
    status = f"Done · **Audio** {duration:.1f}s · **Wall clock** {elapsed:.1f}s · running on CPU"
    if trimmed:
        status += f"\n\n*Input was longer than {MAX_AUDIO_SECONDS:.0f}s and was trimmed.*"
    yield " ".join(pieces).strip(), status


# ---------------------------------------------------------------------------
# Dictation mode: live microphone → real-time streaming transcription
#
# The stock gr.Audio(streaming=True) component only fires its `.stream()`
# callback on fixed chunk boundaries and, in practice, the transcript only
# visibly updated once the user pressed stop. So dictation is now driven by a
# *custom HTML/JS component* (see the DICTATION_HTML / DICTATION_JS strings and
# gr.HTML below): the browser captures mic audio client-side and, several times
# a second, POSTs the newest samples to the backend via Gradio's server_functions
# (the `/component_server` route) and writes the returned partial transcript
# straight into the DOM. The transcript therefore grows chunk-by-chunk *during*
# recording, decoupled from any stop/finalize event.
#
# The backend ASR decoding logic below is unchanged from before — it still
# buffers ~DICTATION_STEP_SECONDS of audio and decodes each segment with the
# same `_asr_engine` call. Only the *delivery* of partial results changed:
# instead of `yield`-ing to a gr.Textbox on stream events, each call returns the
# running transcript to the JS client, which paints it immediately.
# ---------------------------------------------------------------------------

import asyncio
import base64
import threading

# Per-browser-session dictation state, keyed by a client-generated session id.
# Each entry mirrors the old dictation state: a not-yet-decoded audio tail
# (`pending`, float32 @ 24 kHz) and the list of decoded text `pieces`.
_DICTATION_SESSIONS: dict[str, dict] = {}
_DICTATION_LOCK = threading.Lock()


def _new_dictation_state() -> dict:
    """Return a fresh dictation state.

    - ``pending``: not-yet-transcribed audio tail (float32, 24 kHz).
    - ``pieces``: text decoded from each completed segment so far.
    """
    return {"pending": np.zeros(0, dtype=np.float32), "pieces": []}


def _get_session(session_id: str) -> dict:
    """Fetch (creating if needed) the dictation state for a client session."""
    with _DICTATION_LOCK:
        st = _DICTATION_SESSIONS.get(session_id)
        if st is None:
            st = _new_dictation_state()
            _DICTATION_SESSIONS[session_id] = st
            # Bound memory if many sessions accumulate over the life of the app.
            if len(_DICTATION_SESSIONS) > 64:
                for old in list(_DICTATION_SESSIONS)[:-32]:
                    _DICTATION_SESSIONS.pop(old, None)
        return st


def _decode_pending(state: dict, hotwords: str, n_threads: int) -> bool:
    """Decode every full DICTATION_STEP_SECONDS segment in the pending tail.

    Uses exactly the same per-segment `_asr_engine` decode as before, appending
    each segment's text to ``state['pieces']``. Returns True if any new text was
    produced, so the caller knows whether to push a fresh transcript to the UI.
    """
    step = int(DICTATION_STEP_SECONDS * TARGET_SR)
    produced = False
    pending = state["pending"]
    while len(pending) >= step:
        segment, pending = pending[:step], pending[step:]
        text = _asr_engine(segment, hotwords, n_threads, greedy=True)
        if text and text.strip():
            state["pieces"].append(text.strip())
            produced = True
    state["pending"] = pending
    # Cap decoded pieces so very long sessions stay bounded.
    max_pieces = int(DICTATION_MAX_SECONDS / DICTATION_STEP_SECONDS)
    if len(state["pieces"]) > max_pieces:
        state["pieces"] = state["pieces"][-max_pieces:]
    return produced


def dictation_start(data: dict) -> dict:
    """Reset a dictation session when the browser starts a new recording.

    Server function called from the custom HTML component's JS. ``data`` is a
    JSON object with a client-generated ``session_id``.
    """
    session_id = str((data or {}).get("session_id", "default"))
    with _DICTATION_LOCK:
        _DICTATION_SESSIONS[session_id] = _new_dictation_state()
    return {"transcript": "", "status": "Listening… speak into your microphone."}


def _dictation_push_sync(data: dict) -> dict:
    """Blocking core of :func:`dictation_push` (runs the CPU ASR decode).

    Kept as a plain synchronous function so it can be dispatched to a worker
    thread via ``asyncio.to_thread`` — see :func:`dictation_push`. It must not
    run on the event loop directly: each decode is a multi-second blocking
    ``subprocess.run`` and would otherwise stall the whole server, preventing
    partial transcripts from reaching the browser until recording stops.
    """
    data = data or {}
    session_id = str(data.get("session_id", "default"))
    hotwords = str(data.get("hotwords", "") or "")
    n_threads = int(data.get("n_threads", N_CPU) or N_CPU)
    sr = int(data.get("sr", TARGET_SR) or TARGET_SR)

    state = _get_session(session_id)

    b64 = data.get("audio") or ""
    if b64:
        raw = base64.b64decode(b64)
        chunk = np.frombuffer(raw, dtype="<f4").astype(np.float32)
        if chunk.size:
            # Client sends normalised [-1, 1] floats; guard against int16 range.
            peak = float(np.max(np.abs(chunk)))
            if peak > 1.5:
                chunk = chunk / 32768.0
            chunk = _resample_24k(chunk, sr)
            state["pending"] = np.concatenate([state["pending"], chunk])

    step = int(DICTATION_STEP_SECONDS * TARGET_SR)
    if len(state["pending"]) < step:
        return {
            "transcript": " ".join(state["pieces"]).strip(),
            "status": "Listening…",
            "updated": False,
        }

    updated = _decode_pending(state, hotwords, n_threads)
    return {
        "transcript": " ".join(state["pieces"]).strip(),
        "status": "Listening…",
        "updated": updated,
    }


async def dictation_push(data: dict) -> dict:
    """Decode the newest slice of live mic audio and return the running transcript.

    Server function called repeatedly (several times a second) from the custom
    HTML component while the mic is live. Because it returns the partial
    transcript on *every* call, the JS client can paint incremental updates
    continuously during recording, rather than waiting for a stop event.

    Gradio's ``/component_server`` route awaits coroutine server functions on
    the event loop but runs *sync* ones inline on it. The decode is a blocking,
    multi-second ``subprocess.run``; if it ran on the loop it would freeze the
    whole server for its duration, so no earlier partial response could be
    flushed to the browser until recording stopped — which is exactly the bug
    where the transcript only appeared after Stop. Running it in a worker
    thread keeps the loop free to deliver each partial transcript immediately.

    Args:
        data: JSON object with ``session_id`` (str), ``sr`` (int sample rate),
            ``audio`` (base64-encoded little-endian float32 PCM mono samples),
            and optional ``hotwords`` (str) / ``n_threads`` (int).

    Returns:
        JSON object with ``transcript`` (running text), ``status`` (short label)
        and ``updated`` (True when this call produced new decoded text).
    """
    return await asyncio.to_thread(_dictation_push_sync, data)


def _dictation_finalize_sync(data: dict) -> dict:
    """Blocking core of :func:`dictation_finalize` (runs the final ASR decode)."""
    data = data or {}
    session_id = str(data.get("session_id", "default"))
    hotwords = str(data.get("hotwords", "") or "")
    n_threads = int(data.get("n_threads", N_CPU) or N_CPU)

    state = _get_session(session_id)
    pending = state.get("pending")
    if pending is not None and len(pending) > 0:
        text = _asr_engine(pending, hotwords, n_threads, greedy=True)
        if text and text.strip():
            state["pieces"].append(text.strip())
        state["pending"] = np.zeros(0, dtype=np.float32)
    transcript = " ".join(state.get("pieces", [])).strip()
    status = "Done · running on CPU." if transcript else "No audio captured."
    return {"transcript": transcript, "status": status}


async def dictation_finalize(data: dict) -> dict:
    """Decode any leftover audio tail when the browser stops recording.

    Server function called once from the custom HTML component after the user
    presses stop. Runs the blocking flush decode in a worker thread (same
    reasoning as :func:`dictation_push`) and returns the finalised transcript.
    """
    return await asyncio.to_thread(_dictation_finalize_sync, data)


# ---------------------------------------------------------------------------
# Custom HTML/JS dictation component
#
# This is the "custom Gradio HTML component" the maintainer asked for. It does
# NOT use gr.Audio(streaming=True); instead it captures mic audio client-side
# with the Web Audio API and pushes each newly recorded slice to the backend
# via Gradio server_functions (dictation_start / dictation_push /
# dictation_finalize), painting the returned partial transcript into the DOM as
# soon as it arrives — so the transcript grows continuously *during* recording.
# ---------------------------------------------------------------------------

DICTATION_HTML = """
<div class="dictation-widget">
  <div class="dictation-controls">
    <button class="dictation-btn dictation-start" type="button">● Start recording</button>
    <button class="dictation-btn dictation-stop" type="button" disabled>■ Stop</button>
    <span class="dictation-status">Idle — press start and speak.</span>
  </div>
  <label class="dictation-label">Live transcript</label>
  <div class="dictation-transcript" aria-live="polite"></div>
</div>
"""

DICTATION_CSS = """
.dictation-widget { display: flex; flex-direction: column; gap: 8px; }
.dictation-controls { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; }
.dictation-btn {
  padding: 6px 14px; border-radius: 8px; border: 1px solid var(--border-color-primary, #ccc);
  background: var(--button-secondary-background-fill, #f3f4f6);
  color: var(--body-text-color, #111); cursor: pointer; font-size: 14px;
}
.dictation-btn:disabled { opacity: 0.5; cursor: not-allowed; }
.dictation-start { color: #b91c1c; font-weight: 600; }
.dictation-status { font-size: 13px; opacity: 0.85; }
.dictation-label { font-size: 13px; opacity: 0.7; }
.dictation-transcript {
  min-height: 180px; max-height: 420px; overflow-y: auto; padding: 12px;
  border: 1px solid var(--border-color-primary, #ccc); border-radius: 8px;
  background: var(--background-fill-primary, #fff); color: var(--body-text-color, #111);
  white-space: pre-wrap; line-height: 1.5; font-size: 15px;
}
"""

# The client-side driver. `element`, `server`, `props` are provided by Gradio's
# gr.HTML js_on_load runtime. `server.dictation_*` map to the Python functions
# passed via server_functions=[...] and are called over /component_server.
DICTATION_JS = r"""
const el = element;
const startBtn = el.querySelector('.dictation-start');
const stopBtn = el.querySelector('.dictation-stop');
const statusEl = el.querySelector('.dictation-status');
const transcriptEl = el.querySelector('.dictation-transcript');

const sessionId = 'sess-' + Math.random().toString(36).slice(2) + '-' + Date.now();
const TARGET_SR = 24000;
const PUSH_INTERVAL_MS = 700;   // how often we ship the newest audio to the server

let audioCtx = null, source = null, processor = null, stream = null;
let buffer = [];          // accumulated Float32 samples since last push
let recording = false;
let inFlight = false;     // avoid overlapping server calls
let captureSR = 48000;
let pushTimer = null;
// Bumped on every stop(); lets us ignore any live-push response that resolves
// after recording ended so it can't repaint a stale "Listening…" status.
let recordGen = 0;

function setStatus(t) { statusEl.textContent = t; }
function setTranscript(t) {
  transcriptEl.textContent = t || '';
  transcriptEl.scrollTop = transcriptEl.scrollHeight;
}

function floatsToBase64(floats) {
  const bytes = new Uint8Array(floats.buffer, floats.byteOffset, floats.byteLength);
  let bin = '';
  const CH = 0x8000;
  for (let i = 0; i < bytes.length; i += CH) {
    bin += String.fromCharCode.apply(null, bytes.subarray(i, i + CH));
  }
  return btoa(bin);
}

async function pushBuffer(isFinal) {
  if (inFlight) return;
  const gen = recordGen;   // which recording session this push belongs to
  let samples;
  if (buffer.length === 0) {
    if (!isFinal) return;
    samples = new Float32Array(0);
  } else {
    let total = 0;
    for (const b of buffer) total += b.length;
    samples = new Float32Array(total);
    let off = 0;
    for (const b of buffer) { samples.set(b, off); off += b.length; }
    buffer = [];
  }
  inFlight = true;
  try {
    const payload = {
      session_id: sessionId,
      sr: captureSR,
      audio: samples.length ? floatsToBase64(samples) : ''
    };
    const res = isFinal ? await server.dictation_finalize(payload)
                        : await server.dictation_push(payload);
    if (res) {
      // The transcript is always safe to paint. But a *live* push (not the
      // final flush) that resolves after stop() must NOT overwrite the
      // terminal status, or the "Listening…" label lingers after Stop.
      if (typeof res.transcript === 'string') setTranscript(res.transcript);
      if (typeof res.status === 'string' && (isFinal || gen === recordGen)) {
        setStatus(res.status);
      }
    }
  } catch (e) {
    console.warn('dictation push failed', e);
    if (isFinal || gen === recordGen) setStatus('Error talking to server (see console).');
  } finally {
    inFlight = false;
  }
}

async function start() {
  if (recording) return;
  try {
    stream = await navigator.mediaDevices.getUserMedia({ audio: true });
  } catch (e) {
    setStatus('Microphone permission denied.');
    return;
  }
  audioCtx = new (window.AudioContext || window.webkitAudioContext)();
  // Browsers may create the context in a "suspended" state (autoplay policy);
  // while suspended, onaudioprocess never fires and no audio is captured, so
  // resume it explicitly before wiring up the graph.
  try { if (audioCtx.state === 'suspended') await audioCtx.resume(); } catch (e) {}
  captureSR = audioCtx.sampleRate || 48000;
  source = audioCtx.createMediaStreamSource(stream);
  processor = audioCtx.createScriptProcessor(4096, 1, 1);
  processor.onaudioprocess = (ev) => {
    if (!recording) return;
    const input = ev.inputBuffer.getChannelData(0);
    buffer.push(new Float32Array(input));
  };
  source.connect(processor);
  processor.connect(audioCtx.destination);

  recording = true;
  recordGen++;
  buffer = [];
  startBtn.disabled = true;
  stopBtn.disabled = false;
  setTranscript('');
  setStatus('Listening…');
  try { await server.dictation_start({ session_id: sessionId }); } catch (e) {}

  // Continuously ship the newest audio and repaint the transcript while live.
  pushTimer = setInterval(() => { pushBuffer(false); }, PUSH_INTERVAL_MS);
}

async function stop() {
  if (!recording) return;
  recording = false;
  recordGen++;   // invalidate any in-flight live push so it can't restore "Listening…"
  startBtn.disabled = false;
  stopBtn.disabled = true;
  if (pushTimer) { clearInterval(pushTimer); pushTimer = null; }
  try { if (processor) processor.disconnect(); } catch (e) {}
  try { if (source) source.disconnect(); } catch (e) {}
  try { if (audioCtx) await audioCtx.close(); } catch (e) {}
  try { if (stream) stream.getTracks().forEach(t => t.stop()); } catch (e) {}
  // Clear the live "Listening…" indicator immediately on Stop.
  setStatus('Finalising…');
  // Wait for any in-flight live push to finish (so its stale response can't
  // land after ours), then flush the tail through finalize. finalize always
  // runs — even if the wait times out — so we never get stuck on "Finalising…".
  let tries = 0;
  while (inFlight && tries < 100) { await new Promise(r => setTimeout(r, 50)); tries++; }
  inFlight = false;   // force the final flush through regardless of a stuck flag
  try {
    await pushBuffer(true);
  } catch (e) {
    console.warn('finalize failed', e);
  }
  // Guarantee a terminal, idle status even if finalize returned nothing usable.
  const s = statusEl.textContent || '';
  if (s === 'Listening…' || s === 'Finalising…' || s === '') {
    setStatus('Done · running on CPU.');
  }
}

startBtn.addEventListener('click', start);
stopBtn.addEventListener('click', stop);
"""


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

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

with gr.Blocks(title="VibeVoice-ASR-BitNet") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """# 🎙️ VibeVoice-ASR-BitNet
Multilingual speech recognition with a **ternary (1.58-bit) BitNet decoder**, running
entirely on CPU through the authors' ggml engine.

Supports English, Chinese, French, Italian, Korean, Portuguese and Vietnamese.
"""
        )

        with gr.Tabs():
            with gr.Tab("Transcribe file / recording"):
                with gr.Row():
                    with gr.Column():
                        audio_in = gr.Audio(
                            label="Speech",
                            sources=["upload", "microphone"],
                            type="filepath",
                            format="wav",
                        )
                        run_btn = gr.Button("Transcribe", variant="primary")
                    with gr.Column():
                        transcript_out = gr.Textbox(
                            label="Transcript",
                            lines=9,
                            max_lines=25,
                            buttons=["copy"],
                            interactive=False,
                        )
                        metrics_out = gr.Markdown()

                with gr.Accordion("Advanced settings", open=False):
                    hotwords_in = gr.Textbox(
                        label="Hotwords / context",
                        placeholder="VibeVoice, BitNet, ggml",
                        info="Comma-separated terms that bias the decoder towards rare words.",
                        value="",
                    )
                    threads_in = gr.Slider(
                        label="CPU threads",
                        minimum=1, maximum=MAX_THREADS, step=1, value=N_CPU,
                    )
                    greedy_in = gr.Checkbox(
                        label="Greedy decoding",
                        value=True,
                        info="Deterministic and recommended. Uncheck to sample.",
                    )
                    temperature_in = gr.Slider(
                        label="Temperature", minimum=0.1, maximum=1.5, step=0.05, value=0.7,
                    )
                    top_p_in = gr.Slider(
                        label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=0.9,
                    )

                inputs = [audio_in, hotwords_in, threads_in, greedy_in, temperature_in, top_p_in]
                outputs = [transcript_out, metrics_out]

                gr.Examples(
                    examples=[
                        ["examples/librispeech_en_1.wav"],
                        ["examples/librispeech_en_2.wav"],
                        ["examples/fleurs_fr.wav"],
                        ["examples/fleurs_zh.wav"],
                    ],
                    inputs=[audio_in],
                    outputs=outputs,
                    fn=transcribe,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Examples (LibriSpeech · FLEURS fr · FLEURS zh)",
                )

            with gr.Tab("Dictation mode"):
                gr.Markdown(
                    "Speak into your microphone and watch the transcription appear live, "
                    "streamed on CPU as you talk — the text grows chunk-by-chunk **while "
                    "you are still recording**, not only when you stop. Press stop to "
                    "finalise."
                )
                # Custom HTML component: client-side mic capture + continuous
                # partial-transcript delivery via server_functions. No stock
                # gr.Audio(streaming=True) / gr.Textbox event model here.
                dictation_widget = gr.HTML(
                    value="",
                    html_template=DICTATION_HTML,
                    css_template=DICTATION_CSS,
                    js_on_load=DICTATION_JS,
                    server_functions=[
                        dictation_start,
                        dictation_push,
                        dictation_finalize,
                    ],
                    elem_id="dictation-widget",
                    container=True,
                    padding=True,
                )

        gr.Markdown(
            "Model: [microsoft/VibeVoice-ASR-BitNet](https://huggingface.co/microsoft/VibeVoice-ASR-BitNet) · "
            "Engine: [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp) · "
            "Paper: [VibeVoice-ASR-BitNet Technical Report](https://huggingface.co/papers/2607.21075)"
        )

    run_btn.click(fn=transcribe, inputs=inputs, outputs=outputs, api_name="transcribe")

    # Dictation is driven entirely by the custom HTML component's JS talking to
    # the dictation_* server_functions, so there are no gr.Audio stream/stop
    # event listeners to wire up here.

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