#!/usr/bin/env python3
"""Persistent CUDA/TensorRT YAMNet worker for Omnius audio classification.

The process deliberately accepts only already-captured WAV paths.  It never
opens an ALSA device and never downloads a model.  Omnius bootstrap stages and
checksums the ONNX model and TensorRT plan before this worker is launched.

Protocol: JSON Lines on stdin/stdout.  stdout is protocol-only; diagnostics go
to stderr so callers cannot confuse logs with inference results.
"""
from __future__ import annotations

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


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 cuda_check(result, label: str):
    """cuda-python returns either an error enum or (error enum, value)."""
    if isinstance(result, tuple):
        status, value = result[0], result[1:]
    else:
        status, value = result, ()
    if int(status) != 0:
        raise RuntimeError(f"CUDA {label} failed: {status}")
    if len(value) == 0:
        return None
    return value[0] if len(value) == 1 else value


class CudaPythonRuntime:
    """Thin adapter over JetPack's cuda-python bindings."""
    def __init__(self, cudart):
        self.cudart = cudart

    def stream_create(self):
        return cuda_check(self.cudart.cudaStreamCreate(), "stream create")

    def alloc(self, size, label):
        return cuda_check(self.cudart.cudaMalloc(size), f"{label} allocation")

    def h2d(self, device, host, stream, label):
        cuda_check(
            self.cudart.cudaMemcpyAsync(
                device,
                host.ctypes.data,
                host.nbytes,
                self.cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
                stream,
            ),
            f"{label} copy",
        )

    def d2h(self, host, device, stream, label):
        cuda_check(
            self.cudart.cudaMemcpyAsync(
                host.ctypes.data,
                device,
                host.nbytes,
                self.cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
                stream,
            ),
            f"{label} copy",
        )

    def synchronize(self, stream):
        cuda_check(self.cudart.cudaStreamSynchronize(stream), "stream synchronize")

    def free(self, allocation):
        self.cudart.cudaFree(allocation)

    def device_metadata(self):
        device = cuda_check(self.cudart.cudaGetDevice(), "get device")
        props = cuda_check(self.cudart.cudaGetDeviceProperties(device), "get device properties")
        name = getattr(props, "name", "CUDA device")
        if isinstance(name, bytes):
            name = name.split(b"\0", 1)[0].decode("utf8", "replace")
        return str(name), f"{getattr(props, 'major', 0)}.{getattr(props, 'minor', 0)}"


class PyCudaRuntime:
    """JetPack also ships python3-pycuda on supported L4T images.

    Keeping this fallback avoids a generic PyPI CUDA runtime wheel: both
    backends call the CUDA libraries supplied by the installed JetPack image.
    """
    def __init__(self, cuda):
        cuda.init()
        self.cuda = cuda
        self.device = cuda.Device(0)
        self.context = self.device.make_context()

    def stream_create(self):
        return self.cuda.Stream()

    def alloc(self, size, _label):
        return self.cuda.mem_alloc(size)

    def h2d(self, device, host, stream, _label):
        self.cuda.memcpy_htod_async(device, host, stream)

    def d2h(self, host, device, stream, _label):
        self.cuda.memcpy_dtoh_async(host, device, stream)

    def synchronize(self, stream):
        stream.synchronize()

    def free(self, allocation):
        allocation.free()

    def stream_handle(self, stream):
        return stream.handle

    def device_metadata(self):
        major, minor = self.device.compute_capability()
        return self.device.name(), f"{major}.{minor}"


def read_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 != 16000:
        raise RuntimeError(f"Expected 16 kHz WAV from the caller, received {sample_rate} Hz")
    waveform = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    if waveform.size > 96000:
        raise RuntimeError(
            f"Expected at most 6 seconds of caller-conditioned audio, received {waveform.size / sample_rate:.3f} seconds"
        )
    return (
        waveform,
        sample_rate,
        (time.perf_counter() - started) * 1000,
        "sha256:" + hashlib.sha256(raw).hexdigest(),
    )


class YAMNetTensorRT:
    def __init__(self, engine_path: str, class_map_path: str):
        started = time.perf_counter()
        import numpy as np
        import tensorrt as trt
        try:
            from cuda import cudart
            cuda = CudaPythonRuntime(cudart)
        except ImportError:
            import pycuda.driver as pycuda
            cuda = PyCudaRuntime(pycuda)

        self.np = np
        self.trt = trt
        self.cuda = cuda
        self.logger = trt.Logger(trt.Logger.ERROR)
        with open(engine_path, "rb") as handle:
            serialized = handle.read()
        self.runtime = trt.Runtime(self.logger)
        self.engine = self.runtime.deserialize_cuda_engine(serialized)
        if self.engine is None:
            raise RuntimeError("TensorRT could not deserialize the YAMNet engine")
        self.context = self.engine.create_execution_context()
        if self.context is None:
            raise RuntimeError("TensorRT could not create a YAMNet execution context")
        self.stream = cuda.stream_create()
        self.input_name = None
        self.output_names = []
        for index in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(index)
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.input_name = name
            else:
                self.output_names.append(name)
        if self.input_name is None or not self.output_names:
            raise RuntimeError("YAMNet TensorRT engine did not expose input and output tensors")
        self.score_output_name = self._unique_output_name(521, "AudioSet score")
        self.embedding_output_name = self._unique_output_name(1024, "YAMNet embedding")
        with open(class_map_path, newline="", encoding="utf8") as handle:
            self.classes = [row["display_name"] for row in csv.DictReader(handle)]
        if len(self.classes) != 521:
            raise RuntimeError(f"YAMNet class map must contain 521 classes; found {len(self.classes)}")
        self.device_name, self.compute_capability = self._device_metadata()
        self.model_load_ms = (time.perf_counter() - started) * 1000

    def _device_metadata(self):
        # TensorRT uses the CUDA-visible ordinal.  Omnius preflight limits that
        # namespace to exactly the approved Jetson GPU before process launch.
        return self.cuda.device_metadata()

    def _unique_output_name(self, width: int, label: str):
        matches = []
        for name in self.output_names:
            shape = tuple(self.engine.get_tensor_shape(name))
            if shape and int(shape[-1]) == width:
                matches.append(name)
        if len(matches) != 1:
            raise RuntimeError(
                f"YAMNet TensorRT engine must expose exactly one {label} output with width {width}; found {matches}"
            )
        return matches[0]

    def warm(self):
        # YAMNet needs >=0.975 seconds.  This is a model-load probe, not a
        # captured-audio classification, and it verifies the TensorRT path is
        # actually executable before readiness is reported.
        # Exercise and validate both exported model heads before reporting the
        # persistent worker as ready.  A stale score-only plan must fail here,
        # rather than after an embedding caller has acquired the queue slot.
        self.infer(self.np.zeros(15600, dtype=self.np.float32), 1, include_embedding=True)

    def infer(self, waveform, top_k: int | None = None, include_embedding: bool = False):
        started = time.perf_counter()
        np = self.np
        if waveform.ndim != 1 or waveform.size < 15600:
            # YAMNet's framing kernel requires 0.975 s; pad silence only to
            # satisfy the model shape, without altering any supplied samples.
            waveform = np.pad(waveform.reshape(-1), (0, max(0, 15600 - waveform.size)))
        waveform = np.ascontiguousarray(waveform.astype(np.float32, copy=False))
        if not self.context.set_input_shape(self.input_name, tuple(waveform.shape)):
            raise RuntimeError(f"TensorRT rejected YAMNet input shape {tuple(waveform.shape)}")

        allocations = []
        host_outputs = {}
        try:
            input_bytes = waveform.nbytes
            input_device = self.cuda.alloc(input_bytes, "input")
            allocations.append(input_device)
            self.context.set_tensor_address(self.input_name, int(input_device))
            for name in self.output_names:
                shape = tuple(self.context.get_tensor_shape(name))
                if any(int(dim) < 0 for dim in shape):
                    raise RuntimeError(f"TensorRT did not resolve output shape for {name}: {shape}")
                dtype = self.trt.nptype(self.engine.get_tensor_dtype(name))
                host = np.empty(shape, dtype=dtype)
                device = self.cuda.alloc(host.nbytes, name)
                allocations.append(device)
                host_outputs[name] = (host, device)
                self.context.set_tensor_address(name, int(device))
            self.cuda.h2d(input_device, waveform, self.stream, "input")
            stream_handle = self.cuda.stream_handle(self.stream) if hasattr(self.cuda, "stream_handle") else self.stream
            if not self.context.execute_async_v3(stream_handle):
                raise RuntimeError("TensorRT YAMNet execute_async_v3 returned false")
            for name, (host, device) in host_outputs.items():
                self.cuda.d2h(host, device, self.stream, name)
            self.cuda.synchronize(self.stream)
            result = {}
            if top_k is not None:
                scores = host_outputs[self.score_output_name][0]
                if scores.ndim != 2 or scores.shape[-1] != 521:
                    raise RuntimeError(f"YAMNet score output shape is invalid: {scores.shape}")
                mean_scores = np.mean(scores, axis=0)
                indices = np.argsort(mean_scores)[-top_k:][::-1]
                result["classifications"] = [
                    {"label": self.classes[int(index)], "confidence": float(mean_scores[int(index)])}
                    for index in indices
                ]
            if include_embedding:
                frame_embeddings = host_outputs[self.embedding_output_name][0]
                if frame_embeddings.ndim != 2 or frame_embeddings.shape[-1] != 1024 or frame_embeddings.shape[0] < 1:
                    raise RuntimeError(f"YAMNet embedding output shape is invalid: {frame_embeddings.shape}")
                vector = np.mean(frame_embeddings.astype(np.float32, copy=False), axis=0)
                norm = float(np.linalg.norm(vector))
                if not np.isfinite(norm) or norm <= 0.0:
                    raise RuntimeError("YAMNet embedding has an invalid L2 norm")
                result["embedding"] = {
                    "dimension": 1024,
                    "frame_count": int(frame_embeddings.shape[0]),
                    "pooling": "mean_l2",
                    "dtype": "float32",
                    "normalization": "l2",
                    "values": (vector / norm).astype(np.float32).tolist(),
                }
            return result, (time.perf_counter() - started) * 1000
        finally:
            for allocation in allocations:
                try:
                    self.cuda.free(allocation)
                except Exception:
                    pass


def file_digest(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return "sha256:" + digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--engine", required=True)
    parser.add_argument("--class-map", required=True)
    parser.add_argument("--model-digest", required=True)
    args = parser.parse_args()
    # CUDA_VISIBLE_DEVICES must be a single approved accelerator before any
    # CUDA/TensorRT import. Jetson's integrated GPU is logical device 0.
    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
    if visible != "0":
        fail("Audio TensorRT worker requires exactly CUDA_VISIBLE_DEVICES=0 on Jetson")
    worker = YAMNetTensorRT(args.engine, args.class_map)
    worker.warm()
    emit({
        "type": "ready",
        "pid": os.getpid(),
        "backend": "tensorrt-fp16",
        "device": worker.device_name,
        "compute_capability": worker.compute_capability,
        "cuda_visible_devices": visible,
        "model": "yamnet",
        "model_digest": args.model_digest,
        "taxonomy": "AudioSet-521",
        "model_load_ms": round(worker.model_load_ms, 3),
        "warmed": True,
    })
    for line in sys.stdin:
        try:
            request = json.loads(line)
            if request.get("type") == "shutdown":
                emit({"type": "stopped"})
                return 0
            action = str(request.get("action") or request.get("type") or "").strip().lower()
            if action not in ("classify", "embed"):
                raise RuntimeError("unknown audio analysis 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("audio analysis requires id and file")
            waveform, sample_rate, decode_ms, input_digest = read_wav(file_path, worker.np)
            rms = float(worker.np.sqrt(worker.np.mean(worker.np.square(waveform)))) if waveform.size else 0.0
            peak = float(worker.np.max(worker.np.abs(waveform))) if waveform.size else 0.0
            # Do not force YAMNet to invent an AudioSet label for digital or
            # near-digital silence. This remains a valid, grounded result.
            low_information = rms < 0.0005 and peak < 0.002
            if low_information:
                result, inference_ms = (
                    {"classifications": []} if action == "classify" else {"embedding": None},
                    0.0,
                )
            else:
                result, inference_ms = worker.infer(
                    waveform,
                    top_k=max(1, min(int(request.get("top_k", 5)), 25)) if action == "classify" else None,
                    include_embedding=action == "embed",
                )
            emit({
                "type": "result",
                "id": request_id,
                "action": action,
                "success": True,
                "sample_rate_hz": sample_rate,
                "duration_seconds": round(float(waveform.size) / sample_rate, 6),
                "input_digest": input_digest,
                "low_information": low_information,
                "acoustic": {"rms": round(rms, 8), "peak": round(peak, 8)},
                "timings_ms": {"decode": round(decode_ms, 3), "inference": round(inference_ms, 3)},
                **result,
            })
        except Exception as exc:  # keep the persistent worker alive per request
            emit({
                "type": "result",
                "id": str(locals().get("request", {}).get("id", "")),
                "action": str(locals().get("action", "")),
                "success": False,
                "error": str(exc),
            })
            print(traceback.format_exc(), file=sys.stderr, flush=True)
    return 0


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