#!/usr/bin/env python3
"""Persistent, offline CLAP semantic-audio embedding worker.

This process deliberately has no setup behavior.  The Node runtime has already
created the system-site venv, installed its pinned non-Torch requirements, and
verified the immutable model artifact before this process is started.
"""

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

if os.environ.get("HF_HUB_OFFLINE") != "1" or os.environ.get("TRANSFORMERS_OFFLINE") != "1":
    raise RuntimeError("CLAP worker must be started with offline model access enabled")
if os.environ.get("PYTHONNOUSERSITE") != "1":
    raise RuntimeError("CLAP worker must run from the isolated managed venv without user-site packages")

import numpy as np
import torch
import torch.nn.functional as F
from transformers import ClapModel, ClapProcessor

MODEL_ID = "laion/clap-htsat-unfused"
# The worker receives exactly Egg's caller-conditioned audio window.  It does
# not decode arbitrary media, downmix, trim, or otherwise silently redefine
# the observation that will later be persisted.
SOURCE_SAMPLE_RATE_HZ = 16000
MAX_SOURCE_SAMPLES = 96000
SAMPLE_RATE_HZ = 48000
MAX_SAMPLES = 288000
EMBEDDING_DIMENSION = 512


def json_line(value):
    print(json.dumps(value, separators=(",", ":")), flush=True)


def sha256_file(path):
    digest = hashlib.sha256()
    with open(path, "rb") as source:
        while True:
            block = source.read(1024 * 1024)
            if not block:
                break
            digest.update(block)
    return "sha256:" + digest.hexdigest()


def read_pcm_wav(path):
    with wave.open(path, "rb") as handle:
        if handle.getcomptype() != "NONE":
            raise ValueError("Expected caller-conditioned uncompressed mono PCM16/16 kHz WAV")
        channels = handle.getnchannels()
        width = handle.getsampwidth()
        sample_rate = handle.getframerate()
        frames = handle.readframes(handle.getnframes())
    if channels != 1 or width != 2 or sample_rate != SOURCE_SAMPLE_RATE_HZ:
        raise ValueError("Expected caller-conditioned mono PCM16/16 kHz WAV")
    if not frames or len(frames) % 2:
        raise ValueError("WAV contains no complete PCM16 samples")
    if len(frames) // 2 > MAX_SOURCE_SAMPLES:
        raise ValueError("Expected caller-conditioned WAV of at most 6 seconds")
    # The digest is the exact PCM window actually consumed, not mutable RIFF
    # metadata such as timestamp, comment, or chunk ordering.
    input_digest = "sha256:" + hashlib.sha256(frames).hexdigest()
    return np.frombuffer(frames, dtype="<i2").astype(np.float32) / 32768.0, input_digest


def resample_conditioned_16k_to_48k(samples):
    # SOURCE_SAMPLE_RATE_HZ divides CLAP's rate exactly. Linear interpolation
    # keeps preprocessing explicit and deterministic without adding another
    # resampling dependency to JetPack's pinned venv.
    target_count = samples.size * 3
    old = np.arange(samples.size, dtype=np.float64)
    new = np.linspace(0, max(0, samples.size - 1), target_count, dtype=np.float64)
    return np.interp(new, old, samples).astype(np.float32)


class ClapSemanticWorker:
    def __init__(self, model_dir, model_digest, model_revision):
        actual = sha256_file(os.path.join(model_dir, "pytorch_model.bin"))
        if actual != model_digest:
            raise RuntimeError("CLAP weight digest mismatch: expected %s, got %s" % (model_digest, actual))
        self.model_digest = model_digest
        self.model_revision = model_revision
        self.device = torch.device("cuda:0")
        if not torch.cuda.is_available():
            raise RuntimeError("JetPack CUDA Torch is unavailable")
        if torch.cuda.get_device_capability(0) != (8, 7) or not str(torch.version.cuda or "").startswith("12.2"):
            raise RuntimeError("CLAP requires the preflighted JetPack Orin CUDA:0 (CC 8.7 / CUDA 12.2)")
        self.processor = ClapProcessor.from_pretrained(model_dir, local_files_only=True)
        self.model = ClapModel.from_pretrained(model_dir, local_files_only=True).to(self.device).eval()
        if getattr(self.model.config, "projection_dim", None) != EMBEDDING_DIMENSION:
            raise RuntimeError("Pinned CLAP model did not expose a 512-dimensional projection")
        self._warm()

    def _features(self, samples):
        encoded = self.processor(audios=samples, sampling_rate=SAMPLE_RATE_HZ, return_tensors="pt")
        encoded = {key: value.to(self.device) for key, value in encoded.items() if isinstance(value, torch.Tensor)}
        with torch.inference_mode():
            output = self.model.get_audio_features(**encoded)
            output = F.normalize(output.float(), p=2, dim=-1)
        vector = output[0].detach().cpu().numpy().astype(np.float32)
        if vector.shape != (EMBEDDING_DIMENSION,) or not np.isfinite(vector).all():
            raise RuntimeError("CLAP did not return one finite 512-dimensional vector")
        return vector

    def _warm(self):
        self._features(np.zeros(MAX_SAMPLES, dtype=np.float32))
        torch.cuda.synchronize(self.device)

    def embed(self, path):
        started = time.perf_counter()
        samples, input_digest = read_pcm_wav(path)
        decode_ms = (time.perf_counter() - started) * 1000
        source_duration_seconds = samples.size / SOURCE_SAMPLE_RATE_HZ
        samples = resample_conditioned_16k_to_48k(samples)
        rms = float(np.sqrt(np.mean(np.square(samples))))
        peak = float(np.max(np.abs(samples)))
        duration_seconds = source_duration_seconds
        common = {
            "sample_rate_hz": SOURCE_SAMPLE_RATE_HZ,
            "semantic_sample_rate_hz": SAMPLE_RATE_HZ,
            "duration_seconds": duration_seconds,
            "input_digest": input_digest,
            "acoustic": {"rms": rms, "peak": peak},
            "timings_ms": {"decode": decode_ms, "inference": 0.0},
        }
        if rms < 0.0005 and peak < 0.002:
            return dict(common, low_information=True, embedding=None, segment_count=0)
        inference_started = time.perf_counter()
        vectors = []
        # The contract bounds inputs to six seconds, so one explicit 48 kHz
        # window is embedded; semantic segment pooling is kept in the schema
        # for model-set compatibility, not used to hide arbitrary long input.
        vectors.append(self._features(samples))
        vector = np.mean(np.stack(vectors), axis=0)
        vector /= max(float(np.linalg.norm(vector)), 1e-12)
        inference_ms = (time.perf_counter() - inference_started) * 1000
        return dict(common,
            low_information=False,
            segment_count=len(vectors),
            embedding={"dimension": EMBEDDING_DIMENSION, "pooling": "segment_mean_l2", "dtype": "float32", "normalization": "l2", "values": vector.tolist()},
            timings_ms={"decode": decode_ms, "inference": inference_ms},
        )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--model-digest", required=True)
    parser.add_argument("--model-revision", required=True)
    args = parser.parse_args()
    worker = ClapSemanticWorker(args.model_dir, args.model_digest, args.model_revision)
    capability = torch.cuda.get_device_capability(0)
    json_line({"type": "ready", "pid": os.getpid(), "backend": "transformers-clap", "device": torch.cuda.get_device_name(0), "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), "compute_capability": "%d.%d" % capability, "torch_cuda_version": str(torch.version.cuda), "model": MODEL_ID, "model_revision": args.model_revision, "model_digest": args.model_digest, "embedding_dimension": EMBEDDING_DIMENSION, "sample_rate_hz": SAMPLE_RATE_HZ, "warmed": True})
    for line in sys.stdin:
        try:
            request = json.loads(line)
            if request.get("type") != "embed" or not isinstance(request.get("id"), str) or not isinstance(request.get("file"), str):
                raise ValueError("Expected an embed request with id and file")
            result = worker.embed(request["file"])
            json_line(dict(result, type="result", id=request["id"], action="embed", success=True))
        except Exception as error:
            json_line({"type": "result", "id": request.get("id") if "request" in locals() and isinstance(request, dict) else "", "action": "embed", "success": False, "error": str(error)})


if __name__ == "__main__":
    main()