#!/usr/bin/env python3
"""
live-nemotron.py — Self-contained streaming ASR worker using NVIDIA's
nvidia/nemotron-speech-streaming-en-0.6b model.

Parallel to live-whisper.py. Same stdin/stdout protocol so the same
pipelines (nexus voice subsystem, asr_listen tool, eval harness) can
swap backends by pointing at a different script.

Protocol:
  stdin   — raw PCM16 (16kHz, mono, 16-bit signed little-endian)
  stdout  — JSON lines:
    {"type":"status","message":"Creating venv..."}
    {"type":"status","message":"Installing dependencies..."}
    {"type":"status","message":"Loading model..."}
    {"type":"ready"}
    {"type":"transcript","text":"hello world","isFinal":false}
    {"type":"transcript","text":"hello world how are you","isFinal":true}
    {"type":"error","message":"..."}

Usage:
  # Live stream from mic:
  arecord -f S16_LE -r 16000 -c 1 -t raw -q - | python3 live-nemotron.py
  # Single file transcription (write path + read transcript JSON):
  python3 live-nemotron.py --file recording.wav

Backend selection:
  1. NeMo toolkit (nvidia NeMo) — native streaming support for Parakeet-
     style models. Preferred when available.
  2. transformers + torchaudio — fallback via HuggingFace's generic
     ASR pipeline. Works for file-based transcription even when NeMo
     install fails (common on macOS / no-CUDA setups). Does NOT do
     streaming — buffers the full window each chunk.
"""

import sys
import os
import json
import subprocess
import struct
import time
import threading
import argparse
from pathlib import Path

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

SCRIPT_DIR = Path(__file__).resolve().parent
VENV = SCRIPT_DIR / ".nemotron-venv"
PY = VENV / "bin" / "python"
PIP = VENV / "bin" / "pip"

SAMPLE_RATE = 16000
CHANNELS = 1
SAMPLE_WIDTH = 2  # 16-bit
CHUNK_SECONDS = 2.0    # Nemotron is a streaming model — shorter chunks than whisper
WINDOW_SECONDS = 8.0

# HuggingFace model identifier
MODEL_ID = "nvidia/nemotron-speech-streaming-en-0.6b"

# All managed ASR weights live outside the npm package so upgrades do not
# invalidate an already-pulled model. HuggingFace honors HF_HOME for both the
# NeMo and transformers loaders below.
MODEL_CACHE_ROOT = Path(
    os.environ.get(
        "OMNIUS_ASR_MODEL_DIR",
        str(Path.home() / ".omnius" / "models" / "asr" / "nemotron-streaming"),
    )
).expanduser()
os.environ.setdefault("HF_HOME", str(MODEL_CACHE_ROOT))

# ---------------------------------------------------------------------------
# Output helpers (JSON lines to stdout)
# ---------------------------------------------------------------------------

def emit(event: dict):
    sys.stdout.write(json.dumps(event) + "\n")
    sys.stdout.flush()


def emit_status(msg: str):
    emit({"type": "status", "message": msg})


def emit_error(msg: str):
    emit({"type": "error", "message": msg})


def emit_transcript(text: str, is_final: bool = False, backend: str = "nemotron"):
    emit({"type": "transcript", "text": text, "isFinal": is_final, "backend": backend})


def asr_cpu_allowed() -> bool:
    return os.environ.get("OMNIUS_ASR_ALLOW_CPU", "").strip().lower() in ("1", "true", "yes", "on")


def select_asr_device() -> str:
    """Select CUDA whenever it is present; never silently retry on CPU."""
    import torch

    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        raw_index = os.environ.get("OMNIUS_ASR_CUDA_DEVICE", "0").strip() or "0"
        try:
            index = int(raw_index)
        except ValueError as error:
            raise RuntimeError(f"Invalid OMNIUS_ASR_CUDA_DEVICE={raw_index!r}") from error
        if index < 0 or index >= torch.cuda.device_count():
            raise RuntimeError(f"OMNIUS_ASR_CUDA_DEVICE={index} is outside CUDA device range 0..{torch.cuda.device_count() - 1}")
        torch.cuda.set_device(index)
        return f"cuda:{index}"
    if asr_cpu_allowed():
        return "cpu"
    cuda_version = getattr(getattr(torch, "version", None), "cuda", None)
    raise RuntimeError(
        f"CUDA-only ASR is enabled but PyTorch cannot use CUDA (torch.version.cuda={cuda_version}, "
        f"device_count={torch.cuda.device_count()}). Set OMNIUS_ASR_ALLOW_CPU=1 only for an explicit emergency fallback."
    )

# ---------------------------------------------------------------------------
# Venv bootstrap (same pattern as live-whisper.py)
# ---------------------------------------------------------------------------

def _in_venv() -> bool:
    # Omnius launches this worker from its managed CUDA Python environment.
    # Requiring the script directory here used to create a second, CPU-prone
    # venv beside the npm package on every Jetson install.
    return sys.prefix != sys.base_prefix


def _ensure_venv():
    if VENV.exists():
        return
    emit_status("Creating Python venv for Nemotron ASR...")
    import venv
    venv.EnvBuilder(with_pip=True).create(str(VENV))
    subprocess.check_call(
        [str(PY), "-m", "pip", "install", "--upgrade", "pip", "wheel"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )


def _ensure_deps():
    """Install torch + either nemo_toolkit[asr] or transformers as fallback."""
    need = []
    try:
        import numpy  # noqa: F401
    except ImportError:
        need.append("numpy")
    try:
        import torch  # noqa: F401
    except ImportError:
        need.append("torch")
    try:
        import soundfile  # noqa: F401
    except ImportError:
        need.append("soundfile")
    try:
        import transformers  # noqa: F401
    except ImportError:
        need.append("transformers")

    if need:
        emit_status(f"Installing core deps: {', '.join(need)}...")
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", *need],
                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
            )
        except subprocess.CalledProcessError as e:
            emit_error(f"pip install failed: {e}")
            sys.exit(1)
        # Force reimport
        for mod in ("numpy", "torch", "soundfile", "transformers"):
            if mod in sys.modules:
                del sys.modules[mod]

    # NeMo toolkit is large and optional — try to install it but fall
    # back gracefully if it's unavailable on this platform.
    try:
        import nemo.collections.asr  # noqa: F401
    except ImportError:
        emit_status("Installing nemo_toolkit[asr] (large — may take a few minutes)...")
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", "nemo_toolkit[asr]"],
                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                timeout=600,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
            emit_status(f"NeMo install skipped ({e}) — will use transformers fallback")

# ---------------------------------------------------------------------------
# Bootstrap: re-exec inside venv
# ---------------------------------------------------------------------------

# --check short-circuit — runs on the host Python without any venv or
# dependency install so CI and smoke tests can verify the script parses
# + is callable without triggering a 5-minute NeMo download.
if "--check" in sys.argv:
    emit({"type": "check", "ok": True, "script": str(Path(__file__).resolve())})
    sys.exit(0)

if not _in_venv():
    _ensure_venv()
    os.execv(str(PY), [str(PY)] + sys.argv)

_ensure_deps()

# Now safe to import
import numpy as np  # noqa: E402

# ---------------------------------------------------------------------------
# Backend loaders
# ---------------------------------------------------------------------------

def _load_nemo_model(model_id: str = MODEL_ID):
    """Load NeMo on the selected CUDA device, with no implicit CPU retry."""
    try:
        import nemo.collections.asr as nemo_asr
        import torch
    except ImportError:
        return (None, None)
    try:
        device = select_asr_device()
        emit_status(f"Loading NeMo model {model_id}...")
        model = nemo_asr.models.ASRModel.from_pretrained(model_id)
        model.eval()
        model = model.to(torch.device(device))
        return (model, device)
    except Exception as e:
        emit_status(f"NeMo load failed: {str(e)[:200]}")
        return (None, None)


def _load_transformers_model(model_id: str = MODEL_ID):
    """Fallback: load via HuggingFace transformers pipeline."""
    try:
        from transformers import pipeline
    except ImportError:
        return None
    try:
        emit_status(f"Loading transformers pipeline for {model_id}...")
        selected = select_asr_device()
        device = int(selected.split(":", 1)[1]) if selected.startswith("cuda:") else -1
        pipe = pipeline(
            task="automatic-speech-recognition",
            model=model_id,
            device=device,
            return_timestamps=False,
            chunk_length_s=30,
            stride_length_s=5,
        )
        return pipe
    except Exception as e:
        emit_status(f"transformers load failed: {e}")
        return None


def _extract_hypothesis_text(r0) -> str:
    """Extract the transcript string from a NeMo result item. Handles
    plain strings, Hypothesis objects (with possibly empty text), and
    nested lists of Hypotheses returned by RNNT models. Returns an
    empty string for silent input rather than dumping the repr."""
    if r0 is None:
        return ""
    if isinstance(r0, str):
        return r0.strip()
    # Nested list of hypotheses (some RNNT decoders)
    if isinstance(r0, list):
        if not r0:
            return ""
        return _extract_hypothesis_text(r0[0])
    # Hypothesis object — may have text="" for silent audio, which is
    # a VALID transcript (just empty). Return it without falling through
    # to str(r0) which would dump the whole repr.
    if hasattr(r0, "text"):
        return str(r0.text or "").strip()
    # best_hypothesis() method (rare)
    if hasattr(r0, "best_hypothesis"):
        try:
            bh = r0.best_hypothesis()
            if bh and hasattr(bh, "text"):
                return str(bh.text or "").strip()
        except Exception:
            pass
    return ""


def _transcribe_buffer_nemo(model, audio: np.ndarray) -> str:
    """Transcribe a 16kHz mono float32 numpy array via NeMo.

    Tries multiple invocation signatures across NeMo versions:
      - transcribe([np.ndarray]) — newest
      - transcribe(paths2audio_files=["file.wav"]) — legacy, requires tmp wav
    """
    try:
        # Newest API: pass audio arrays directly
        result = model.transcribe([audio], batch_size=1, verbose=False)
        if not result:
            return ""
        return _extract_hypothesis_text(result[0])
    except TypeError:
        # Older NeMo — fall through to file path invocation
        pass
    except Exception as e:
        emit_error(f"NeMo transcribe error: {e}")
        return ""

    # Fallback: write audio to a temp WAV and pass the path
    try:
        import soundfile as sf
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            sf.write(tmp.name, audio, SAMPLE_RATE, subtype="PCM_16")
            tmp_path = tmp.name
        try:
            result = model.transcribe(paths2audio_files=[tmp_path], batch_size=1, verbose=False)
            if result and result[0] is not None:
                return _extract_hypothesis_text(result[0])
        finally:
            try:
                os.unlink(tmp_path)
            except Exception:
                pass
    except Exception as e:
        emit_error(f"NeMo legacy transcribe error: {e}")
    return ""


def _transcribe_buffer_transformers(pipe, audio: np.ndarray) -> str:
    """Transcribe via transformers pipeline."""
    try:
        result = pipe({"array": audio, "sampling_rate": SAMPLE_RATE})
        if isinstance(result, dict):
            return str(result.get("text", "")).strip()
        if isinstance(result, list) and result:
            return str(result[0].get("text", "")).strip() if isinstance(result[0], dict) else ""
        return ""
    except Exception as e:
        emit_error(f"transformers transcribe error: {e}")
        return ""

# ---------------------------------------------------------------------------
# File transcription mode (single-shot)
# ---------------------------------------------------------------------------

def transcribe_file(path: str, language: str = "en", model_id: str = MODEL_ID) -> int:
    """Single-file transcription — reads a WAV, prints one transcript
    JSON line, exits. Used by AsrListenTool's file path. Exit code 0
    on success, 1 on failure."""
    try:
        import soundfile as sf
        audio, sr = sf.read(path, dtype="float32")
        if audio.ndim > 1:
            audio = audio.mean(axis=1)  # downmix to mono
        if sr != SAMPLE_RATE:
            # Resample via simple linear interpolation (avoids scipy dep)
            ratio = SAMPLE_RATE / sr
            new_len = int(len(audio) * ratio)
            idx = np.linspace(0, len(audio) - 1, new_len).astype(np.float32)
            audio = np.interp(idx, np.arange(len(audio), dtype=np.float32), audio).astype(np.float32)
    except Exception as e:
        emit_error(f"Failed to load audio file {path}: {e}")
        return 1

    (model, device) = _load_nemo_model(model_id)
    backend = "nemo"
    if model is None:
        device = select_asr_device()
        model = _load_transformers_model(model_id)
        backend = "transformers"
    if model is None:
        emit_error("No nemotron backend available (tried NeMo + transformers)")
        return 1

    emit({"type": "ready", "backend": backend, "device": device or "cpu", "cuda": str(device).startswith("cuda")})

    t0 = time.time()
    if backend == "nemo":
        text = _transcribe_buffer_nemo(model, audio)
    else:
        text = _transcribe_buffer_transformers(model, audio)
    elapsed = time.time() - t0

    # Silent / no-speech audio is NOT an error — it's a valid transcript
    # (empty string). Emit the full envelope so the caller can distinguish
    # "no speech" from "engine crashed". Exit 0 either way.
    emit({
        "type": "transcript",
        "text": text or "",
        "isFinal": True,
        "backend": f"nemotron-{backend}",
        "latencyMs": int(elapsed * 1000),
        "audioSeconds": float(len(audio) / SAMPLE_RATE),
        "empty": not bool(text),
    })
    return 0

# ---------------------------------------------------------------------------
# Streaming mode (stdin → transcripts)
# ---------------------------------------------------------------------------

def stream_stdin(args) -> int:
    (model, _device) = _load_nemo_model(args.model)
    backend = "nemo"
    if model is None:
        model = _load_transformers_model(args.model)
        backend = "transformers"
    if model is None:
        emit_error("No nemotron backend available (tried NeMo + transformers)")
        return 1

    emit({"type": "ready"})

    audio_buf = np.zeros(0, dtype=np.float32)
    buf_lock = threading.Lock()
    chunk_bytes = int(args.chunk_seconds * SAMPLE_RATE * SAMPLE_WIDTH)
    window_samples = int(args.window_seconds * SAMPLE_RATE)
    last_text = ""
    running = True

    def read_stdin():
        nonlocal audio_buf, running
        try:
            while running:
                data = sys.stdin.buffer.read(chunk_bytes)
                if not data:
                    break
                samples = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
                with buf_lock:
                    audio_buf = np.concatenate([audio_buf, samples])
        except Exception:
            pass
        finally:
            running = False

    reader = threading.Thread(target=read_stdin, daemon=True)
    reader.start()

    try:
        while running:
            time.sleep(args.chunk_seconds)
            with buf_lock:
                if len(audio_buf) < SAMPLE_RATE:
                    continue
                window = audio_buf[-window_samples:].copy() if len(audio_buf) > window_samples else audio_buf.copy()
            if backend == "nemo":
                text = _transcribe_buffer_nemo(model, window)
            else:
                text = _transcribe_buffer_transformers(model, window)
            if text and text != last_text:
                last_text = text
                emit_transcript(text, is_final=False, backend=f"nemotron-{backend}")
    except KeyboardInterrupt:
        pass

    with buf_lock:
        full_audio = audio_buf.copy()
    if len(full_audio) >= SAMPLE_RATE:
        if backend == "nemo":
            text = _transcribe_buffer_nemo(model, full_audio)
        else:
            text = _transcribe_buffer_transformers(model, full_audio)
        if text:
            emit_transcript(text, is_final=True, backend=f"nemotron-{backend}")
    running = False
    return 0

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Nemotron streaming ASR worker")
    parser.add_argument("--model", default=MODEL_ID, help="HuggingFace model id (default: nvidia/nemotron-speech-streaming-en-0.6b)")
    parser.add_argument("--file", default=None, help="Transcribe a single audio file instead of stdin")
    parser.add_argument("--language", default="en", help="Language code")
    parser.add_argument("--chunk-seconds", type=float, default=CHUNK_SECONDS, help="Transcribe interval")
    parser.add_argument("--window-seconds", type=float, default=WINDOW_SECONDS, help="Sliding window size")
    parser.add_argument("--stdin", action="store_true", help="Explicit stdin mode (default when no --file)")
    parser.add_argument("--check", action="store_true", help="Just verify the script parses + imports; no model load")
    parser.add_argument("--setup", action="store_true", help="Pull and validate the selected model on the selected device, then exit")
    args = parser.parse_args()

    if args.check:
        emit({"type": "check", "ok": True, "script": str(Path(__file__).resolve())})
        return 0

    if args.setup:
        model, device = _load_nemo_model(args.model)
        backend = "nemo"
        if model is None:
            device = select_asr_device()
            model = _load_transformers_model(args.model)
            backend = "transformers"
        if model is None:
            emit_error("No CUDA Nemotron backend available (tried NeMo + transformers)")
            return 1
        emit({"type": "ready", "backend": backend, "device": device, "cuda": str(device).startswith("cuda"), "model": args.model})
        return 0

    if args.file:
        return transcribe_file(args.file, args.language, args.model)
    return stream_stdin(args)


if __name__ == "__main__":
    sys.exit(main())