#!/usr/bin/env python3
"""
openWakeWord sidecar for voice-surface server.

Protocol (stdio):
  stdin:  raw signed 16-bit PCM at 16kHz, little-endian, mono
          fed in 1280-sample chunks (80ms) by the Node.js parent
  stdout: line-based JSON  {"detected": true, "score": 0.92}  or  {"detected": false, "score": 0.1}

Wake word priority:
  1. Custom sklearn model at models/james_classifier.pkl  (trained with train_james.py)
  2. Fallback: hey_jarvis_v0.1.onnx bundled with openWakeWord
"""
import os, sys, json, pathlib, warnings
warnings.filterwarnings("ignore")
os.environ["ORT_LOGGING_LEVEL"] = "3"

import numpy as np
import openwakeword
from openwakeword.model import Model
from openwakeword.utils import AudioFeatures

SAMPLE_RATE   = 16000
CHUNK_SAMPLES = 1280          # 80 ms per chunk
CHUNK_BYTES   = CHUNK_SAMPLES * 2  # int16 = 2 bytes

THRESHOLD     = float(os.environ.get("OWW_THRESHOLD", "0.5"))
MODELS_DIR    = pathlib.Path(__file__).parent / "models"
JAMES_CLF     = MODELS_DIR / "james_classifier.pkl"

_oww_base     = pathlib.Path(openwakeword.__file__).parent
JARVIS_MODEL  = str(_oww_base / "resources" / "models" / "hey_jarvis_v0.1.onnx")

# ---- Load models ----
use_custom = JAMES_CLF.exists()

if use_custom:
    import joblib
    clf = joblib.load(JAMES_CLF)
    af  = AudioFeatures()
    sys.stdout.write(json.dumps({"ready": True, "model": "james-custom"}) + "\n")
else:
    oww = Model(wakeword_model_paths=[JARVIS_MODEL])
    model_key = list(oww.models.keys())[0]
    sys.stdout.write(json.dumps({"ready": True, "model": "hey_jarvis_fallback"}) + "\n")

sys.stdout.flush()

# Accumulator: build up enough audio to run embed_clips (needs >1280 samples)
# For the custom model we collect 5 chunks (400ms) before scoring to get stable embeddings
# Compute the required detection window size so the embedding model has enough frames.
# openWakeWord's embedding model expects at least ~76 frames. Each frame ≈160 samples.
MIN_EMB_FRAMES = 76
MIN_EMB_SAMPLES = (MIN_EMB_FRAMES + 3) * 160  # mirrors openwakeword utils calculation
CUSTOM_WINDOW = int(np.ceil(MIN_EMB_SAMPLES / CHUNK_SAMPLES))
# Ensure a sensible lower bound (keep previous behaviour if model changes)
CUSTOM_WINDOW = max(5, CUSTOM_WINDOW)
custom_buf = []     # list of int16 arrays

buf = b""

try:
    stdin_raw = sys.stdin.buffer
    while True:
        needed = CHUNK_BYTES - len(buf)
        chunk  = stdin_raw.read(needed)
        if not chunk:
            break
        buf += chunk
        if len(buf) < CHUNK_BYTES:
            continue

        audio_int16 = np.frombuffer(buf[:CHUNK_BYTES], dtype=np.int16)
        buf = buf[CHUNK_BYTES:]

        if use_custom:
            custom_buf.append(audio_int16)
            if len(custom_buf) < CUSTOM_WINDOW:
                # Not enough audio yet — emit low score placeholder
                sys.stdout.write(json.dumps({"detected": False, "score": 0.0}) + "\n")
                sys.stdout.flush()
                continue

            # Got CUSTOM_WINDOW chunks — extract embedding and classify
            audio_window = np.concatenate(custom_buf)
            custom_buf = custom_buf[1:]  # slide window by 1 chunk

            # embed_clips expects a numpy array with shape (N, samples).
            # passing a Python list caused an AttributeError in openwakeword.utils (list has no .shape).
            embs = af.embed_clips(np.stack([audio_window]))   # shape (N, 96)
            if embs is None or len(embs) == 0:
                sys.stdout.write(json.dumps({"detected": False, "score": 0.0}) + "\n")
                sys.stdout.flush()
                continue

            feat = embs.mean(axis=0).reshape(1, -1)
            score = float(clf.predict_proba(feat)[0, 1])
            detected = score >= THRESHOLD

            if detected:
                custom_buf = []   # reset window after detection

        else:
            # hey_jarvis ONNX model — feed 1 chunk at a time
            scores   = oww.predict(audio_int16)
            score    = float(scores.get(model_key, 0.0))
            detected = score >= THRESHOLD
            if detected:
                oww.reset()

        sys.stdout.write(json.dumps({"detected": detected, "score": round(score, 3)}) + "\n")
        sys.stdout.flush()

except (BrokenPipeError, KeyboardInterrupt):
    pass