#!/usr/bin/env python3
"""Persistent CPU-only WeSpeaker CAM++ speaker-identity embedding worker.

This process accepts only caller-owned, already-conditioned mono PCM16/16 kHz
WAV paths.  It never records audio, installs dependencies, downloads models,
or enables CUDA.  The parent runtime checksums the pinned ONNX artifact before
this worker can start.

The fbank path is a pinned, CPU-only NumPy implementation of the Kaldi
configuration used by WeSpeaker's ONNX example: 80 bins, 25 ms / 10 ms Hamming
frames, dither=0 and utterance CMN (without CVN). It intentionally imports
neither Torch nor Torchaudio, so an upstream Torchaudio binary can never bind
against or replace JetPack's vendor Torch ABI.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import time
import traceback
import wave


# Set before importing ONNX Runtime so a globally configured Jetson CUDA device
# cannot become an accidental provider for speaker embeddings.
os.environ["CUDA_VISIBLE_DEVICES"] = ""

MODEL_NAME = "wespeaker-voxceleb-campplus"
SAMPLE_RATE = 16000
EMBEDDING_DIMENSION = 512
MIN_DURATION_SECONDS = 1.5
MAX_DURATION_SECONDS = 30.0
MIN_RMS = 0.003
MIN_PEAK = 0.01
MAX_CLIPPED_FRACTION = 0.01
FBANK_IMPLEMENTATION = "kaldi-fbank-numpy-v1"
# A CPU FFT may differ in harmless last bits across x86/aarch64 and NumPy
# builds. Validate the public Kaldi configuration with numeric invariants,
# not a byte hash of implementation-dependent FFT output.
FBANK_VALIDATION = "kaldi-fbank-invariants-v2"
CALIBRATION_BAND_MEANS = (11.1405, 17.5231, 12.8845, 9.2557, 7.7213, 6.7842, 6.6405)


def emit(payload: dict) -> None:
    sys.stdout.write(json.dumps(payload, separators=(",", ":")) + "\n")
    sys.stdout.flush()


def fail(message: str) -> None:
    emit({"type": "fatal", "error": message})
    raise RuntimeError(message)


def read_pcm16_16khz_wav(path: str, np):
    started = time.perf_counter()
    with wave.open(path, "rb") as reader:
        channels = reader.getnchannels()
        sample_width = reader.getsampwidth()
        sample_rate = reader.getframerate()
        frames = reader.getnframes()
        compression = reader.getcomptype()
        raw = reader.readframes(frames)
    if compression != "NONE":
        raise RuntimeError("Only uncompressed RIFF/WAV input is supported")
    if channels != 1:
        raise RuntimeError(f"Expected mono WAV from the caller, received {channels} channels")
    if sample_width != 2:
        raise RuntimeError(f"Expected PCM16 WAV from the caller, received {sample_width * 8}-bit samples")
    if sample_rate != SAMPLE_RATE:
        raise RuntimeError(f"Expected 16 kHz WAV from the caller, received {sample_rate} Hz")
    samples = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    if samples.size > int(MAX_DURATION_SECONDS * SAMPLE_RATE):
        raise RuntimeError(
            f"Expected at most {MAX_DURATION_SECONDS:.0f} seconds of caller-conditioned speech, received {samples.size / sample_rate:.3f} seconds"
        )
    return (
        samples,
        sample_rate,
        "sha256:" + hashlib.sha256(raw).hexdigest(),
        (time.perf_counter() - started) * 1000,
    )


def acoustic_metrics(samples, np):
    if samples.size == 0:
        return {"rms": 0.0, "peak": 0.0, "clipped_fraction": 0.0}
    absolute = np.abs(samples)
    return {
        "rms": float(np.sqrt(np.mean(np.square(samples)))),
        "peak": float(np.max(absolute)),
        "clipped_fraction": float(np.mean(absolute >= 0.999)),
    }


def quality_gate(samples, metrics):
    duration = samples.size / SAMPLE_RATE
    if duration < MIN_DURATION_SECONDS:
        return "insufficient_duration"
    if metrics["rms"] < MIN_RMS or metrics["peak"] < MIN_PEAK:
        return "low_energy"
    if metrics["clipped_fraction"] > MAX_CLIPPED_FRACTION:
        return "clipped"
    return "usable"


def kaldi_fbank80_numpy(waveform, np):
    """Kaldi-compatible fbank for the fixed WeSpeaker CAM++ ONNX contract.

    This mirrors WeSpeaker's published Kaldi fbank configuration: mono 16k
    input, snip_edges=True,
    remove_dc_offset=True, preemphasis=0.97, round_to_power_of_two=True,
    80 triangular 20Hz..Nyquist mel bins, Hamming window and log power. The
    caller passes PCM-scaled float32 samples (not normalized [-1, 1] values).
    """
    dtype = np.float32
    waveform = np.ascontiguousarray(waveform, dtype=dtype)
    frame_size = 400
    frame_shift = 160
    fft_size = 512
    if waveform.ndim != 1 or waveform.size < frame_size:
        raise RuntimeError("Kaldi fbank requires at least one 25 ms mono frame")
    frame_count = 1 + (waveform.size - frame_size) // frame_shift
    frames = np.lib.stride_tricks.as_strided(
        waveform,
        shape=(frame_count, frame_size),
        strides=(frame_shift * waveform.strides[0], waveform.strides[0]),
        writeable=False,
    ).copy()
    # Published Kaldi fbank frame/window defaults.
    frames -= np.mean(frames, axis=1, dtype=dtype, keepdims=True)
    frames[:, 1:] -= dtype(0.97) * frames[:, :-1]
    frames[:, 0] -= dtype(0.97) * frames[:, 0]
    sample_index = np.arange(frame_size, dtype=dtype)
    hamming = dtype(0.54) - dtype(0.46) * np.cos(
        dtype(2.0 * np.pi) * sample_index / dtype(frame_size - 1)
    )
    frames *= hamming
    padded = np.pad(frames, ((0, 0), (0, fft_size - frame_size)), mode="constant")
    # NumPy's rfft may calculate internally at float64 precision. Cast its
    # power spectrum back to float32 before the Kaldi-style matrix multiply.
    spectrum = np.asarray(np.abs(np.fft.rfft(padded, axis=1)) ** 2, dtype=dtype)

    nyquist = dtype(8000.0)
    low_hz = dtype(20.0)
    mel = lambda hz: dtype(1127.0) * np.log(dtype(1.0) + hz / dtype(700.0))
    mel_low = mel(low_hz)
    mel_high = mel(nyquist)
    mel_step = (mel_high - mel_low) / dtype(81.0)
    bin_index = np.arange(80, dtype=dtype)[:, None]
    left = mel_low + bin_index * mel_step
    center = mel_low + (bin_index + dtype(1.0)) * mel_step
    right = mel_low + (bin_index + dtype(2.0)) * mel_step
    fft_bin_hz = dtype(16000.0 / fft_size) * np.arange(fft_size // 2, dtype=dtype)[None, :]
    fft_bin_mel = mel(fft_bin_hz)
    up = (fft_bin_mel - left) / (center - left)
    down = (right - fft_bin_mel) / (right - center)
    banks = np.maximum(dtype(0.0), np.minimum(up, down))
    banks = np.pad(banks, ((0, 0), (0, 1)), mode="constant")
    energies = np.matmul(spectrum, np.asarray(banks.T, dtype=dtype))
    return np.log(np.maximum(energies, np.finfo(dtype).eps)).astype(dtype, copy=False)


def validate_kaldi_fbank_numpy(np):
    """Validate the no-Torch preprocessing path before readiness is true.

    The calibration intentionally asserts shape, finite/CMN invariants, the
    expected 220/440 Hz Kaldi mel-band profile and energy range with tolerances
    that are stable across supported CPU FFT implementations. It must reject a
    changed window, frame layout, PCM scale, mel bank, pre-emphasis or CMN, but
    must not reject Jetson due to last-bit numerical differences.
    """
    t = np.arange(SAMPLE_RATE * 2, dtype=np.float32) / np.float32(SAMPLE_RATE)
    samples = (
        np.float32(0.075) * np.sin(np.float32(2.0 * np.pi * 220.0) * t)
        + np.float32(0.025) * np.sin(np.float32(2.0 * np.pi * 440.0) * t)
    ).astype(np.float32)
    raw_features = kaldi_fbank80_numpy(samples * np.float32(1 << 15), np)
    features = (
        raw_features - np.mean(raw_features, axis=0, dtype=np.float32, keepdims=True)
    ).astype(np.float32)
    if features.shape != (198, 80) or not np.all(np.isfinite(features)):
        raise RuntimeError("NumPy Kaldi fbank calibration produced an invalid feature matrix")
    cmn_abs_mean_max = float(np.max(np.abs(np.mean(features, axis=0))))
    if cmn_abs_mean_max > 1e-3:
        raise RuntimeError("NumPy Kaldi fbank calibration failed utterance CMN")
    raw_min = float(np.min(raw_features))
    raw_max = float(np.max(raw_features))
    cmn_rms = float(np.sqrt(np.mean(np.square(features))))
    band_means = raw_features.mean(axis=0)[[0, 5, 10, 20, 40, 60, 79]]
    if not (4.5 < raw_min < 6.0 and 19.0 < raw_max < 21.5 and 0.55 < cmn_rms < 0.75):
        raise RuntimeError(
            "NumPy Kaldi fbank calibration energy/range invariant failed"
        )
    if not np.allclose(band_means, np.asarray(CALIBRATION_BAND_MEANS, dtype=np.float32), rtol=0.0, atol=0.25):
        raise RuntimeError("NumPy Kaldi fbank calibration mel-band invariant failed")
    return {
        "validation": FBANK_VALIDATION,
        "frames": int(features.shape[0]),
        "bins": int(features.shape[1]),
        "cmn_abs_mean_max": round(cmn_abs_mean_max, 8),
        "cmn_rms": round(cmn_rms, 6),
        "raw_feature_range": [round(raw_min, 6), round(raw_max, 6)],
    }


class WeSpeakerCamPlus:
    def __init__(self, model_path: str):
        started = time.perf_counter()
        import numpy as np
        import onnxruntime as ort

        self.np = np
        self.preprocessing_validation = validate_kaldi_fbank_numpy(np)
        options = ort.SessionOptions()
        options.inter_op_num_threads = 1
        options.intra_op_num_threads = 1
        self.session = ort.InferenceSession(
            model_path,
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        providers = self.session.get_providers()
        if providers != ["CPUExecutionProvider"]:
            raise RuntimeError(f"WeSpeaker must run only on CPUExecutionProvider, got {providers}")
        inputs = self.session.get_inputs()
        outputs = self.session.get_outputs()
        if len(inputs) != 1 or inputs[0].name != "feats" or inputs[0].type != "tensor(float)":
            raise RuntimeError("Pinned CAM++ ONNX model must expose float input 'feats'")
        if len(inputs[0].shape) != 3 or inputs[0].shape[-1] != 80:
            raise RuntimeError(f"Pinned CAM++ input must be [batch, frames, 80], got {inputs[0].shape}")
        if len(outputs) != 1 or outputs[0].name != "embs" or outputs[0].type != "tensor(float)":
            raise RuntimeError("Pinned CAM++ ONNX model must expose float output 'embs'")
        if len(outputs[0].shape) != 2 or outputs[0].shape[-1] != EMBEDDING_DIMENSION:
            raise RuntimeError(f"Pinned CAM++ output must be [batch, {EMBEDDING_DIMENSION}], got {outputs[0].shape}")
        self.model_load_ms = (time.perf_counter() - started) * 1000

    def features(self, samples):
        # WeSpeaker's ONNX inference multiplies normalized PCM values by 2**15
        # before Kaldi fbank. Decode follows the same convention above.
        mat = kaldi_fbank80_numpy(samples * self.np.float32(1 << 15), self.np)
        if mat.ndim != 2 or mat.shape[0] < 1 or mat.shape[1] != 80:
            raise RuntimeError(f"NumPy Kaldi fbank produced an invalid shape: {tuple(mat.shape)}")
        # CMN without CVN, over the full input utterance, per official script.
        mat = mat - self.np.mean(mat, axis=0, dtype=self.np.float32, keepdims=True)
        return mat[None, :, :].astype(self.np.float32, copy=False)

    def embed(self, samples):
        started = time.perf_counter()
        features = self.features(samples)
        output = self.session.run(["embs"], {"feats": features})[0]
        vector = self.np.asarray(output, dtype=self.np.float32)
        if vector.shape != (1, EMBEDDING_DIMENSION):
            raise RuntimeError(f"CAM++ produced invalid embedding shape: {vector.shape}")
        vector = vector[0]
        norm = float(self.np.linalg.norm(vector))
        if not self.np.isfinite(norm) or norm <= 0.0:
            raise RuntimeError("CAM++ embedding has an invalid L2 norm")
        return {
            "dimension": EMBEDDING_DIMENSION,
            "dtype": "float32",
            "normalization": "l2",
            "values": (vector / norm).astype(self.np.float32).tolist(),
        }, (time.perf_counter() - started) * 1000

    def warm(self):
        # A deterministic voiced-like probe validates the official fbank and
        # the CPU session before readiness is published. It is not user audio.
        t = self.np.arange(SAMPLE_RATE * 2, dtype=self.np.float32) / SAMPLE_RATE
        samples = (0.1 * self.np.sin(2.0 * self.np.pi * 220.0 * t)).astype(self.np.float32)
        embedding, _ = self.embed(samples)
        if len(embedding["values"]) != EMBEDDING_DIMENSION:
            raise RuntimeError("CAM++ warmup did not produce a 512-dimensional embedding")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model")
    parser.add_argument("--model-digest")
    parser.add_argument("--preprocessing-probe", action="store_true")
    args = parser.parse_args()
    if args.preprocessing_probe:
        import numpy as np
        emit(
            {
                "type": "preprocessing_probe",
                "implementation": FBANK_IMPLEMENTATION,
                "validation": validate_kaldi_fbank_numpy(np),
            }
        )
        return 0
    if not args.model or not args.model_digest:
        parser.error("--model and --model-digest are required unless --preprocessing-probe is used")
    worker = WeSpeakerCamPlus(args.model)
    worker.warm()
    emit(
        {
            "type": "ready",
            "pid": os.getpid(),
            "backend": "onnxruntime-cpu",
            "provider": "CPUExecutionProvider",
            "model": MODEL_NAME,
            "model_digest": args.model_digest,
            "dimension": EMBEDDING_DIMENSION,
            "model_load_ms": round(worker.model_load_ms, 3),
            "warmed": True,
            "preprocessing": FBANK_IMPLEMENTATION,
            "preprocessing_validation": worker.preprocessing_validation,
        }
    )
    for line in sys.stdin:
        request = None
        action = ""
        try:
            request = json.loads(line)
            if request.get("type") == "shutdown":
                emit({"type": "stopped"})
                return 0
            action = str(request.get("type") or "").strip().lower()
            if action != "embed":
                raise RuntimeError("unknown speaker embedding action")
            request_id = str(request.get("id") or "")
            file_path = str(request.get("file") or "")
            if not request_id or not file_path:
                raise RuntimeError("speaker embedding requires id and file")
            samples, sample_rate, input_digest, decode_ms = read_pcm16_16khz_wav(file_path, worker.np)
            metrics = acoustic_metrics(samples, worker.np)
            quality = quality_gate(samples, metrics)
            if quality == "usable":
                embedding, inference_ms = worker.embed(samples)
            else:
                embedding, inference_ms = None, 0.0
            emit(
                {
                    "type": "result",
                    "id": request_id,
                    "action": "embed",
                    "success": True,
                    "sample_rate_hz": sample_rate,
                    "input_digest": input_digest,
                    "duration_seconds": round(float(samples.size) / sample_rate, 6),
                    "quality": quality,
                    "embedding": embedding,
                    "acoustic": {key: round(value, 8) for key, value in metrics.items()},
                    "timings_ms": {"decode": round(decode_ms, 3), "inference": round(inference_ms, 3)},
                }
            )
        except Exception as exc:  # Keep the persistent worker alive per request.
            emit(
                {
                    "type": "result",
                    "id": str(request.get("id", "")) if isinstance(request, dict) else "",
                    "action": action,
                    "success": False,
                    "error": str(exc),
                }
            )
            print(traceback.format_exc(), file=sys.stderr, flush=True)
    return 0


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