#!/usr/bin/env python3
"""Offline, line-delimited workers for Omnius managed speaker diarization.

The parent process supplies an already-verified local artifact and forces
offline Hugging Face mode.  This worker deliberately has no setup/download
branch and never reads a token: a request can only use a model proven during
the admin setup path.
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
import wave
from pathlib import Path
from typing import Any


def emit(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, separators=(",", ":")), flush=True)


def validate_live_wav(path: str) -> None:
    """Live Sortformer is a bounded retained 16k mono PCM16 window."""
    try:
        with wave.open(path, "rb") as wav:
            channels = wav.getnchannels()
            width = wav.getsampwidth()
            rate = wav.getframerate()
            duration = wav.getnframes() / float(rate) if rate else 0.0
    except Exception as exc:  # wave errors are intentionally surfaced to caller
        raise ValueError(f"live diarization requires readable WAV: {exc}") from exc
    if channels != 1 or width != 2 or rate != 16000:
        raise ValueError("live diarization requires retained mono PCM16/16 kHz WAV")
    if duration <= 0 or duration > 6.0:
        raise ValueError("live diarization accepts retained audio windows greater than 0 and at most 6 seconds")


def load_pcm16_waveform(path: str):
    """Decode locally without TorchCodec, which has no CPython 3.10/aarch64 wheel."""
    import numpy as np
    import torch

    try:
        with wave.open(path, "rb") as wav:
            if wav.getcomptype() != "NONE" or wav.getsampwidth() != 2:
                raise ValueError("offline reconciliation requires uncompressed PCM16 WAV")
            channels = wav.getnchannels()
            sample_rate = wav.getframerate()
            frames = wav.readframes(wav.getnframes())
    except Exception as exc:
        raise ValueError(f"offline reconciliation requires readable PCM16 WAV: {exc}") from exc
    if channels < 1 or sample_rate < 1 or not frames:
        raise ValueError("offline reconciliation WAV is empty or invalid")
    samples = np.frombuffer(frames, dtype="<i2")
    if samples.size % channels:
        raise ValueError("offline reconciliation WAV has an incomplete PCM frame")
    waveform = samples.reshape(-1, channels).T.astype(np.float32) / 32768.0
    return {
        "waveform": torch.from_numpy(waveform),
        "sample_rate": sample_rate,
        "uri": Path(path).stem,
    }


def number(value: Any) -> float | None:
    try:
        result = float(value)
        return result if result == result and result not in (float("inf"), float("-inf")) else None
    except Exception:
        return None


def append_segment(result: list[dict[str, Any]], start: Any, end: Any, speaker: Any, confidence: Any = None) -> None:
    start_s = number(start)
    end_s = number(end)
    label = str(speaker or "").strip()
    if start_s is None or end_s is None or end_s <= start_s or not label:
        return
    confidence_value = number(confidence)
    result.append({
        "start_ms": round(start_s * 1000),
        "end_ms": round(end_s * 1000),
        "speaker": label,
        "confidence": confidence_value,
    })


def normalize_sortformer(value: Any) -> list[dict[str, Any]]:
    """Accept NeMo's current tuple/list result variants without inventing data."""
    rttm_text = isinstance(value, str)
    if isinstance(value, str):
        value = [line for line in value.splitlines() if line.strip()]
    while isinstance(value, list) and len(value) == 1 and isinstance(value[0], (list, tuple)):
        value = value[0]
    result: list[dict[str, Any]] = []
    if not isinstance(value, (list, tuple)):
        raise RuntimeError("Sortformer returned an unsupported diarization result shape")
    for item in value:
        if isinstance(item, dict):
            append_segment(result, item.get("start", item.get("start_time")), item.get("end", item.get("end_time")), item.get("speaker", item.get("label")), item.get("confidence"))
        elif isinstance(item, (tuple, list)) and len(item) >= 3:
            append_segment(result, item[0], item[1], item[2], item[3] if len(item) > 3 else None)
        elif isinstance(item, str):
            # Some NeMo surfaces return RTTM lines. Parse only the standard
            # SPEAKER fields; malformed lines are ignored rather than guessed.
            fields = item.split()
            if len(fields) >= 8 and fields[0].upper() == "SPEAKER":
                start = number(fields[3])
                duration = number(fields[4])
                if start is not None and duration is not None:
                    append_segment(result, start, start + duration, fields[7])
    # A native RTTM command may emit diagnostics but no SPEAKER rows for a
    # genuine silent window. Structured NeMo objects that are non-empty yet
    # unparseable remain a hard contract failure.
    if not result and value and not rttm_text:
        raise RuntimeError("Sortformer result contained no parseable speaker turns")
    return result


def load_live(args: argparse.Namespace):
    if args.native_binary:
        binary = Path(args.native_binary)
        if not binary.is_file() or not os.access(binary, os.X_OK):
            raise RuntimeError("managed nemo-speech binary is missing or not executable")
        # Metadata inspection validates that this exact Q8 artifact is readable
        # without keeping a second Python/Torch stack on JetPack.
        inspected = subprocess.run(
            [str(binary), "model", "info", args.artifact_path],
            text=True,
            capture_output=True,
            timeout=60,
            check=False,
        )
        if inspected.returncode != 0:
            raise RuntimeError(
                f"nemo-speech rejected the pinned Sortformer model: {(inspected.stderr or inspected.stdout)[-1000:]}"
            )

        def diarize_native(path: str) -> list[dict[str, Any]]:
            validate_live_wav(path)
            completed = subprocess.run(
                [
                    str(binary),
                    "diarize",
                    path,
                    "--model",
                    args.artifact_path,
                    "--device",
                    "cuda:0",
                    "--format",
                    "rttm",
                ],
                text=True,
                capture_output=True,
                timeout=90,
                check=False,
            )
            if completed.returncode != 0:
                raise RuntimeError(
                    f"nemo-speech diarization failed: {(completed.stderr or completed.stdout)[-2000:]}"
                )
            return normalize_sortformer(completed.stdout)

        return binary, diarize_native, "nemo-speech-cpp-q8-cuda"

    import torch
    from nemo.collections.asr.models import SortformerEncLabelModel

    if not torch.cuda.is_available() or tuple(torch.cuda.get_device_capability(0)) != (8, 7):
        raise RuntimeError("live Sortformer requires Jetson Orin CUDA:0 (compute capability 8.7)")
    model = SortformerEncLabelModel.restore_from(
        args.artifact_path, map_location=torch.device("cuda:0"), strict=False
    )
    model.eval()

    def diarize(path: str) -> list[dict[str, Any]]:
        validate_live_wav(path)
        with torch.inference_mode():
            return normalize_sortformer(model.diarize(audio=[path], batch_size=1))

    return model, diarize, "nemo-sortformer-python-cuda"


def load_reconcile(args: argparse.Namespace):
    # Community-1's supported offline workflow is Pipeline.from_pretrained on
    # a local clone.  HF token/terms are setup concerns and intentionally never
    # appear here. CPU protects the live CUDA sidecar on Jetson.
    import torch
    from pyannote.audio import Pipeline

    pipeline = Pipeline.from_pretrained(args.snapshot_path)
    if pipeline is None:
        raise RuntimeError("pyannote failed to load the local Community-1 snapshot")
    pipeline.to(torch.device("cpu"))

    def diarize(path: str) -> list[dict[str, Any]]:
        if not Path(path).is_file():
            raise ValueError("offline reconciliation audio file does not exist")
        # pyannote.audio officially supports in-memory waveform mappings. This
        # avoids its optional TorchCodec decoder, which has no aarch64 wheel,
        # without weakening the model or touching JetPack CUDA Torch.
        output = pipeline(load_pcm16_waveform(path))
        annotation = getattr(output, "speaker_diarization", output)
        turns = getattr(annotation, "itertracks", None)
        if not callable(turns):
            raise RuntimeError("pyannote Community-1 result has no speaker_diarization tracks")
        result: list[dict[str, Any]] = []
        for segment, _track, speaker in annotation.itertracks(yield_label=True):
            append_segment(result, getattr(segment, "start", None), getattr(segment, "end", None), speaker)
        return result

    return pipeline, diarize, "pyannote-community-python-cpu"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["live", "reconcile"], required=True)
    parser.add_argument("--snapshot-path", required=True)
    parser.add_argument("--artifact-path", required=True)
    parser.add_argument("--revision", required=True)
    parser.add_argument("--model-digest", required=True)
    parser.add_argument("--native-binary")
    args = parser.parse_args()

    # Do not allow an accidentally inherited setup token to become a request
    # time credential or network fallback.
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    os.environ.pop("HF_TOKEN", None)
    os.environ.pop("HUGGING_FACE_HUB_TOKEN", None)
    try:
        _model, diarize, backend = load_live(args) if args.mode == "live" else load_reconcile(args)
        model = "nvidia/diar_streaming_sortformer_4spk-v2" if args.mode == "live" else "pyannote/speaker-diarization-community-1"
        emit({
            "type": "ready",
            "pid": os.getpid(),
            "mode": args.mode,
            "model": model,
            "revision": args.revision,
            "model_digest": args.model_digest,
            "backend": backend,
            "warmed": True,
        })
    except Exception as exc:
        emit({"type": "fatal", "error": str(exc)})
        return 2

    for line in sys.stdin:
        try:
            message = json.loads(line)
            if message.get("type") != "diarize" or not isinstance(message.get("id"), str):
                continue
            started = time.perf_counter()
            segments = diarize(str(message.get("file", "")))
            emit({
                "type": "result",
                "id": message["id"],
                "success": True,
                "segments": segments,
                "timings_ms": {"inference": round((time.perf_counter() - started) * 1000)},
            })
        except Exception as exc:
            emit({"type": "result", "id": message.get("id", "unknown") if "message" in locals() else "unknown", "success": False, "error": str(exc)})
    return 0


if __name__ == "__main__":
    raise SystemExit(main())