"""NepaliConformer demo — transcribe Nepali speech, tuned for real telephone audio."""
import os

import gradio as gr
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 nemo.collections.asr.models import EncDecHybridRNNTCTCBPEModel

_path = hf_hub_download("ampixa/nepali-conformer-offline", "nepali_conformer_offline.nemo")
torch.set_num_threads(2)
MODEL = EncDecHybridRNNTCTCBPEModel.restore_from(_path, map_location="cpu")
MODEL.eval()
print("model preloaded", flush=True)

DESCRIPTION = (
    "121M Conformer trained on ~1,655 h of conversational Nepali. "
    "**33.8% WER on real call-center audio** (NepTel benchmark) where Whisper-large-v3 "
    "zero-shot scores ~99%. Honest limitations and the full benchmark: "
    "[github.com/Ampixa/nepaliconformer](https://github.com/Ampixa/nepaliconformer). "
    "CPU demo — a 30 s clip takes roughly 10-20 s. Example clips are real call-center "
    "audio (CC-BY-4.0, © InfoBayAI)."
)

ARTICLE = (
    "Recording transcribes itself as soon as you press ⏹ stop. "
    "Mic blocked? [Open the demo full-screen](https://voidash-nepaliconformer.hf.space) · "
    "Model downloads and usage: "
    "[github.com/Ampixa/nepaliconformer](https://github.com/Ampixa/nepaliconformer#download--run)"
)


def get_model():
    return MODEL


def _prepare(path):
    """Normalize any input to mono 16 kHz PCM and return (path, duration_seconds).

    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.
    """
    import subprocess
    import tempfile

    import soundfile as sf

    try:
        info = sf.info(path)
        needs_convert = info.channels != 1 or info.samplerate != 16000
        duration = info.duration
    except Exception:
        needs_convert = True
        duration = None

    if not needs_convert:
        return path, duration

    out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    proc = subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-i", path,
         "-ar", "16000", "-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]
        )
    return out, sf.info(out).duration


def transcribe(audio_path):
    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 = _prepare(audio_path)
    if duration is not None and duration > 60:
        return "Please keep clips under 60 seconds for this CPU demo."
    model = get_model()
    out = model.transcribe([audio_path], batch_size=1, verbose=False)[0]
    text = out.text if hasattr(out, "text") else str(out)
    return " ".join(t for t in text.split() if t != "<breath>") or "(no speech detected)"


example_files = [[f"examples/{f}"] for f in sorted(os.listdir("examples"))] \
    if os.path.isdir("examples") else []

with gr.Blocks(title="NepaliConformer — Nepali ASR for real telephone calls") as demo:
    gr.Markdown("# NepaliConformer — Nepali ASR for real telephone calls")
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(
                sources=["microphone", "upload"], type="filepath",
                label="Nepali speech (mic or file, ≤60 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, audio_in, text_out)
    audio_in.upload(transcribe, audio_in, text_out)
    submit_btn.click(transcribe, audio_in, text_out)
    clear_btn.click(lambda: (None, ""), None, [audio_in, text_out])

    if example_files:
        gr.Examples(
            examples=example_files, inputs=audio_in, outputs=text_out,
            fn=transcribe, cache_examples=False, run_on_click=True,
        )
    gr.Markdown(ARTICLE)

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