"""Rani demo: transcribe Nepali speech. Ampixa Labs."""
import json
import os
import subprocess
import tempfile

import gradio as gr
import numpy as np
import soundfile as sf
import torch

# Preload at startup: the checkpoint is 485 MB — downloading it lazily made the first
# user's request hang for minutes and look broken.
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import GatedRepoError
from nemo.collections.asr.models import EncDecHybridRNNTCTCBPEModel
from omegaconf import OmegaConf, open_dict

MODEL_REPO = "ampixa/rani-asr"
MAX_SECONDS = 60
SAMPLE_RATE = 16000
# A mic recording that starts with this much exact digital silence lost its first words; one
# that ends with this much was blanked by the browser before upload.
LEAD_SILENCE_NOTE_S = 0.5
TRAIL_SILENCE_NOTE_S = 1.0

# Gradio 5.44's recorder (wavesurfer RecordPlugin) opens the microphone when Record is clicked
# and starts MediaRecorder at once. A microphone delivers exact digital zeros while it starts
# up (measured 1.8 s on Chromium; Bluetooth headsets are typically 1-2 s), so whatever the
# visitor said first was recorded as blanks. Hold getUserMedia until the stream carries real
# samples: the timer and Stop button only appear once recording starts, so they now mean
# "live". Gradio also ignores its own input-device dropdown, so honour it here. Every failure
# path falls back to the unpatched behaviour.
MIC_HEAD = """
<script>
(() => {
  // Brave's fingerprinting protection rescales an AudioBuffer in place on every
  // getChannelData() call (x0.99982 per call, measured on Brave 1.94). Gradio's WAV encoder
  // calls it once per sample, so the uploaded recording decayed to digital zeros after about
  // a second (the speech waveform turned into blanks) and encoding 5 s took ~15 s. Every
  // browser returns the same live array on repeated calls, so caching it per buffer changes
  // nothing elsewhere and stops the compounding. Installed only when reading mutates data.
  if (typeof AudioBuffer === "undefined") return;
  let mutatesOnRead = false;
  try {
    const probe = new AudioBuffer({ length: 128, sampleRate: 8000 });
    probe.getChannelData(0)[0] = 0.5;
    mutatesOnRead = probe.getChannelData(0)[0] !== 0.5;
  } catch (e) {
    console.warn("[audio] getChannelData probe failed; leaving it unpatched", e);
    return;
  }
  if (!mutatesOnRead) return;
  const original = AudioBuffer.prototype.getChannelData;
  const cache = new WeakMap();
  AudioBuffer.prototype.getChannelData = function (channel) {
    let channels = cache.get(this);
    if (!channels) { channels = []; cache.set(this, channels); }
    return channels[channel] || (channels[channel] = original.call(this, channel));
  };
  console.info("[audio] getChannelData mutates data on read (Brave Shields); caching it");
})();
(() => {
  const md = navigator.mediaDevices;
  if (!md || !md.getUserMedia || md.__nepaliConformerPatched) return;
  md.__nepaliConformerPatched = true;
  const original = md.getUserMedia.bind(md);
  const MAX_WAIT_MS = 3000;

  function withSelectedDevice(constraints) {
    const sel = document.querySelector("select.mic-select");
    const id = sel && sel.value;
    if (!id || !constraints || constraints.audio !== true) return null;
    return { ...constraints, audio: { deviceId: { exact: id } } };
  }

  async function waitForSignal(stream) {
    const AC = window.AudioContext || window.webkitAudioContext;
    if (!AC || stream.getAudioTracks().length === 0) return;
    const ctx = new AC();
    try {
      if (ctx.state !== "running") await ctx.resume();
      if (ctx.state !== "running") {
        console.warn("[mic] AudioContext is " + ctx.state + "; not waiting for signal");
        return;
      }
      const source = ctx.createMediaStreamSource(stream);
      const analyser = ctx.createAnalyser();
      analyser.fftSize = 2048;
      const mute = ctx.createGain();
      mute.gain.value = 0;
      source.connect(analyser);
      analyser.connect(mute);
      mute.connect(ctx.destination);
      const buf = new Float32Array(analyser.fftSize);
      const t0 = performance.now();
      while (performance.now() - t0 < MAX_WAIT_MS) {
        analyser.getFloatTimeDomainData(buf);
        if (buf.some((v) => v !== 0)) {
          console.info("[mic] live after " + Math.round(performance.now() - t0) + " ms");
          return;
        }
        await new Promise((r) => setTimeout(r, 20));
      }
      console.warn("[mic] no signal after " + MAX_WAIT_MS + " ms; recording anyway");
    } finally {
      ctx.close().catch((e) => console.warn("[mic] AudioContext close failed", e));
    }
  }

  async function openWarm(constraints) {
    const btn = document.querySelector("button.record-button");
    if (btn) { btn.dataset.warming = "1"; btn.disabled = true; }
    try {
      let stream = null;
      const chosen = withSelectedDevice(constraints);
      if (chosen) {
        try { stream = await original(chosen); }
        catch (e) { console.warn("[mic] selected input device failed; using the default", e); }
      }
      if (!stream) stream = await original(constraints);
      try { await waitForSignal(stream); }
      catch (e) { console.warn("[mic] start-up check failed; recording anyway", e); }
      return stream;
    } finally {
      if (btn) { delete btn.dataset.warming; btn.disabled = false; }
    }
  }

  // Gradio asks for the microphone three times per recording but stops only the last stream,
  // so two stay open and the browser's mic indicator stays on after Stop. Hand out one stream
  // while it is live: Gradio's single stop then closes it.
  let shared = null;
  const isLive = (s) => s.getAudioTracks().some((t) => t.readyState === "live");

  md.getUserMedia = function (constraints) {
    if (!constraints || !constraints.audio || constraints.video) return original(constraints);
    if (shared) {
      const current = shared;
      return current.then((s) => {
        if (isLive(s)) return s;
        if (shared === current) shared = null;
        return md.getUserMedia(constraints);
      });
    }
    const opening = openWarm(constraints);
    shared = opening;
    // A refused permission must not stick: let the next click ask again. The caller still
    // receives the rejection, so Gradio shows its own error.
    opening.catch(() => { if (shared === opening) shared = null; });
    return opening;
  };
})();
</script>
<style>.record-button[data-warming] { opacity: 0.55; cursor: progress; }</style>
"""

try:
    _path = hf_hub_download(MODEL_REPO, "rani/rani.nemo")
except GatedRepoError as e:
    # The model is gated (manual review), so anonymous downloads get a 401. The Space
    # reads its token from the HF_TOKEN secret; without it startup dies in a long
    # traceback that never says what to fix.
    raise RuntimeError(
        f"{MODEL_REPO} is gated and this Space could not authenticate. Set an HF_TOKEN "
        f"secret with read access to {MODEL_REPO} in the Space settings."
    ) from e
torch.set_num_threads(2)
MODEL = EncDecHybridRNNTCTCBPEModel.restore_from(_path, map_location="cpu", strict=False)
MODEL.eval()
# The released checkpoint carries no dataset config; transcribe() needs this field.
with open_dict(MODEL.cfg):
    MODEL.cfg.validation_ds = OmegaConf.create({"use_start_end_token": False})
MODEL.change_decoding_strategy(MODEL.cfg.decoding, decoder_type="rnnt")
print("model preloaded", flush=True)

DESCRIPTION = (
    "Nepali speech-to-text from Ampixa Labs, previously called NepaliConformer. "
    "**33.5 mistakes per 100 words on [NepTel 0.2](https://huggingface.co/datasets/ampixa/neptel)**, "
    "real Nepali phone audio with human-corrected references. "
    "[Results](https://ampixa.com/rani/) · [Code](https://github.com/Ampixa/rani-asr) · "
    "[Weights, with approval](https://huggingface.co/ampixa/rani-asr). "
    "This is a CPU demo, so a 25 s clip takes several seconds."
)

ARTICLE = (
    "**Microphone:** press ⏺ Record, start speaking once the timer runs, press ⏹ Stop. It "
    f"transcribes automatically. Recordings longer than {MAX_SECONDS} s are cut to the first "
    f"{MAX_SECONDS} s. Mic blocked? "
    "[Open the demo full-screen](https://voidash-rani-asr.hf.space).\n\n"
    "**Examples:** phone-call clips are from the "
    "[NepTel](https://huggingface.co/datasets/ampixa/neptel) benchmark (audio © InfoBayAI, "
    "CC-BY-4.0). Read-speech clips are from Google FLEURS ne_np test (CC-BY-4.0). "
    "This demo runs Rani without a language model."
)


def get_model():
    return MODEL


def _prepare(path):
    """Normalize any input to mono 16 kHz PCM capped at MAX_SECONDS.

    Returns (path, original_duration_seconds, temp_files). Browser microphone recordings
    arrive as stereo (and sometimes as webm/mp4 that libsndfile cannot open at all). The
    model takes a mono signal only: feeding it a two-channel file raises "Output shape
    expected = (batch, time)" inside NeMo.
    """
    temps = []
    try:
        info = sf.info(path)
        needs_convert = info.channels != 1 or info.samplerate != SAMPLE_RATE
    except RuntimeError as e:  # soundfile.LibsndfileError: container libsndfile can't read
        print(f"soundfile cannot open {path!r} ({e}); converting with ffmpeg", flush=True)
        needs_convert = True

    if needs_convert:
        out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
        temps.append(out)
        proc = subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", path,
             "-ar", str(SAMPLE_RATE), "-ac", "1", "-c:a", "pcm_s16le", out],
            capture_output=True, text=True,
        )
        if proc.returncode != 0:
            raise gr.Error(
                "Could not read that audio (ffmpeg: %s). Try a WAV or MP3 file."
                % (proc.stderr.strip().splitlines() or ["unknown error"])[-1]
            )
        path = out

    duration = sf.info(path).duration
    if duration > MAX_SECONDS:
        # Cut instead of refusing: a visitor who talked for 70 s should still get a
        # transcript, not an error after the fact.
        data, sr = sf.read(path, frames=MAX_SECONDS * SAMPLE_RATE, dtype="int16")
        out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
        temps.append(out)
        sf.write(out, data, sr, subtype="PCM_16")
        path = out
    return path, duration, temps


def _digital_silence(path):
    """Return (leading_zero_s, trailing_zero_s, nonzero_fraction) of a mono 16 kHz PCM file.

    Exact zeros never come from a live microphone (even a quiet room has noise). Leading zeros
    mean the browser recorded before the device delivered audio; trailing zeros mean the
    browser blanked the audio before upload (Brave's fingerprinting protection did this to
    Gradio's WAV encoder); all zeros mean a dead input.
    """
    data, sr = sf.read(path, dtype="int16")
    if len(data) == 0:
        return 0.0, 0.0, 0.0
    nonzero = np.flatnonzero(data)
    if nonzero.size == 0:
        return len(data) / sr, len(data) / sr, 0.0
    return nonzero[0] / sr, (len(data) - 1 - nonzero[-1]) / sr, nonzero.size / len(data)


def transcribe(audio_path, from_mic=False):
    if not audio_path:
        # Most often: the user pressed Transcribe while the mic was still recording,
        # so no file exists yet.
        return "Press ⏹ stop to finish the recording — it transcribes automatically."
    audio_path, duration, temps = _prepare(audio_path)
    note = ""
    try:
        if from_mic:
            lead, trail, live_fraction = _digital_silence(audio_path)
            print(f"mic recording: {duration:.1f}s, leading digital silence {lead:.2f}s, "
                  f"trailing {trail:.2f}s, non-zero {live_fraction:.0%}", flush=True)
            if live_fraction == 0.0:
                return ("The recording is completely silent (digital zeros): the browser captured "
                        "no sound from the microphone. Check that the browser may use the "
                        "microphone and that the right input is your system default.")
            if lead >= LEAD_SILENCE_NOTE_S:
                note += (f"\n\n(the first {lead:.1f} s of the recording were silent while the "
                         "microphone started; anything said then was not captured)")
            if trail >= TRAIL_SILENCE_NOTE_S:
                note += (f"\n\n(the last {trail:.1f} s of the recording reached the server as "
                         "digital silence: the browser blanked it before upload)")
        model = get_model()
        out = model.transcribe([audio_path], batch_size=1, verbose=False)[0]
    finally:
        for t in temps:
            try:
                os.remove(t)
            except OSError as e:
                print(f"could not remove temp file {t}: {e}", flush=True)
    text = out.text if hasattr(out, "text") else str(out)
    text = " ".join(t for t in text.split() if t != "<breath>") or "(no speech detected)"
    if duration > MAX_SECONDS:
        text += f"\n\n(only the first {MAX_SECONDS} s of {duration:.0f} s were transcribed)"
    return text + note


def transcribe_mic(audio_path):
    return transcribe(audio_path, from_mic=True)


def transcribe_live(audio_path):
    return transcribe(audio_path)


def transcribe_example(audio_path):
    return transcribe(audio_path)


def _load_samples():
    manifest = os.path.join("examples", "samples.json")
    if not os.path.isfile(manifest):
        print(f"no {manifest}; starting without examples", flush=True)
        return []
    with open(manifest, encoding="utf-8") as f:
        samples = json.load(f)
    for s in samples:
        path = os.path.join("examples", s["file"])
        if not os.path.isfile(path):
            raise FileNotFoundError(f"{manifest} lists {s['file']} but {path} is missing")
        s["path"] = path
    return samples


SAMPLES = _load_samples()

with gr.Blocks(title="Rani: Nepali speech-to-text", head=MIC_HEAD) as demo:
    gr.Markdown("# Rani (रानी): Nepali speech-to-text")
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(
                sources=["microphone", "upload"], type="filepath",
                label=f"Nepali speech (mic or file, ≤{MAX_SECONDS} s)",
            )
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Transcribe", variant="primary")
        with gr.Column():
            text_out = gr.Textbox(label="Transcript (Devanagari)", lines=6)

    # Transcribe as soon as the recording stops or a file lands: waiting for an explicit
    # Transcribe click made the demo look broken for anyone who never pressed stop.
    audio_in.stop_recording(transcribe_mic, audio_in, text_out)
    audio_in.upload(transcribe_live, audio_in, text_out)
    submit_btn.click(transcribe, audio_in, text_out)
    clear_btn.click(lambda: (None, ""), None, [audio_in, text_out])

    if SAMPLES:
        gr.Examples(
            examples=[[s["path"]] for s in SAMPLES],
            example_labels=[s["label"] for s in SAMPLES],
            inputs=[audio_in], outputs=text_out,
            fn=transcribe_example, cache_examples=False, run_on_click=True,
            label="Examples — click one to transcribe it",
        )
    gr.Markdown(ARTICLE)

if __name__ == "__main__":
    demo.launch(ssr_mode=False, show_error=True)