#!/usr/bin/env python3
import sys
import os
import signal

# Cleanly handle parent pipe closure on SIGTERM / exit without throwing BrokenPipeError
if hasattr(signal, 'SIGPIPE'):
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)

# Silence OpenCV C++ logging and prioritize DirectShow over MSMF on Windows
os.environ["OPENCV_LOG_LEVEL"] = "SILENT"
os.environ["OPENCV_VIDEOIO_PRIORITY_MSMF"] = "0"
os.environ["OPENCV_FFMPEG_LOGLEVEL"] = "-8"

import json
import time
import collections
import numpy as np
import sounddevice as sd
import cv2
import threading
import base64

SAMPLE_RATE = 48000
BLOCK_SIZE = int(SAMPLE_RATE / 24)  # 24 FPS -> 2000 samples per frame
NUM_MEL_BINS = 40
VISUAL_W = 160
VISUAL_H = 120
prev_audio_fft = None
mel_filterbank_128 = None
hann_window_2000 = None
active_sample_rate = 48000
running = True

def create_mel_filterbank(num_filters=128, n_fft=2000, fs=48000, f_min=20.0, f_max=16000.0):
    """Creates a 128-channel logarithmic/Mel tonotopic cochlear filterbank covering 20Hz-16kHz."""
    def hz_to_mel(hz):
        return 2595.0 * np.log10(1.0 + hz / 700.0)
    def mel_to_hz(mel):
        return 700.0 * (10.0**(mel / 2595.0) - 1.0)

    mel_min = hz_to_mel(f_min)
    mel_max = hz_to_mel(f_max)
    mel_points = np.linspace(mel_min, mel_max, num_filters + 2)
    hz_points = mel_to_hz(mel_points)
    
    n_freqs = n_fft // 2 + 1
    fft_freqs = np.linspace(0, fs / 2, n_freqs)
    
    fb = np.zeros((num_filters, n_freqs), dtype=np.float32)
    for m in range(1, num_filters + 1):
        f_left = hz_points[m - 1]
        f_center = hz_points[m]
        f_right = hz_points[m + 1]
        
        # Rising slope
        up_mask = (fft_freqs >= f_left) & (fft_freqs <= f_center)
        if f_center > f_left:
            fb[m - 1, up_mask] = (fft_freqs[up_mask] - f_left) / (f_center - f_left)
            
        # Falling slope
        down_mask = (fft_freqs >= f_center) & (fft_freqs <= f_right)
        if f_right > f_center:
            fb[m - 1, down_mask] = (f_right - fft_freqs[down_mask]) / (f_right - f_center)
            
    return fb

mel_filterbank_128 = None
hann_window_2000 = None

def extract_f0_pitch(pcm, fs=48000):
    """
    Physical F0 pitch estimation via Normalized Autocorrelation (ACF).
    Extracts fundamental vocal cord frequency in Hz:
    - Adult male: ~85 - 140 Hz
    - Adult female: ~165 - 255 Hz
    - Young child: ~250 - 400 Hz
    """
    n_samples = len(pcm)
    if n_samples < 800:
        return 0.0
    
    # Compute autocorrelation via FFT
    rfft_pcm = np.fft.rfft(pcm, n=n_samples * 2)
    autocorr = np.fft.irfft(rfft_pcm * np.conj(rfft_pcm))[:n_samples]
    
    if autocorr[0] <= 1e-6:
        return 0.0
    norm_ac = autocorr / autocorr[0]
    
    # Search in human vocal range [65 Hz, 500 Hz]
    min_lag = int(fs / 500.0) # 96 samples (~500 Hz)
    max_lag = min(int(fs / 65.0), n_samples - 1) # 738 samples (~65 Hz)
    
    segment = norm_ac[min_lag:max_lag]
    if len(segment) == 0:
        return 0.0
    
    peak_idx = int(np.argmax(segment))
    peak_val = segment[peak_idx]
    
    # Harmonic periodicity threshold
    if peak_val > 0.32:
        best_lag = float(min_lag + peak_idx)
        # Parabolic interpolation for sub-sample precision
        if 0 < peak_idx < len(segment) - 1:
            alpha = float(segment[peak_idx - 1])
            beta = float(segment[peak_idx])
            gamma = float(segment[peak_idx + 1])
            denom = 2.0 * (2.0 * beta - alpha - gamma)
            if abs(denom) > 1e-6:
                delta = (alpha - gamma) / denom
                best_lag += delta
        f0 = fs / max(best_lag, 1.0)
        return float(np.clip(f0, 60.0, 500.0))
    return 0.0

def get_audio_features(indata):
    global prev_audio_fft, mel_filterbank_128, hann_window_2000, active_sample_rate
    if len(indata) == 0:
        return 0, [0]*128, 0, 0.0, False, 0.0, 0.0
    
    fs = active_sample_rate if active_sample_rate else 48000
    pcm = indata[:, 0]
    n_samples = len(pcm)
    if hann_window_2000 is None or len(hann_window_2000) != n_samples:
        hann_window_2000 = np.hanning(n_samples).astype(np.float32)
    if mel_filterbank_128 is None or mel_filterbank_128.shape[1] != (n_samples // 2 + 1):
        mel_filterbank_128 = create_mel_filterbank(128, n_fft=n_samples, fs=fs, f_min=20.0, f_max=min(16000.0, fs / 2.0))

    rms = float(np.sqrt(np.mean(pcm**2)))
    # Continuous full-dynamic range SPL calculation (captures all environmental sounds, clicks, whispers, and speech)
    db = float(np.clip(20.0 * np.log10(rms * 100000.0 + 1e-6) + 15.0, 0.0, 100.0)) if rms > 1e-7 else 0.0
    f0_pitch = extract_f0_pitch(pcm, fs=fs) if rms > 1e-6 else 0.0
    
    # 1. Hann Windowed Real FFT
    windowed_pcm = pcm * hann_window_2000
    fft_vals = np.abs(np.fft.rfft(windowed_pcm)).astype(np.float32)
    
    # 2. 128-Channel Cochlear Tonotopic Filterbank (20Hz to 16kHz)
    tonotopic_bins = np.dot(mel_filterbank_128, fft_vals)
    
    # 3. Peak-relative spectral scaling to preserve formant peaks and spectral contour
    max_val = np.max(tonotopic_bins)
    if max_val > 1e-5:
        norm_fft = (tonotopic_bins / max_val).astype(np.float32)
    else:
        norm_fft = tonotopic_bins.astype(np.float32)
    
    # 4. Calculate spectral flux for transient onset / acoustic burst detection (all sounds)
    flux = 0.0
    if prev_audio_fft is not None:
        diff = norm_fft - prev_audio_fft
        flux = float(np.sum(np.maximum(0.0, diff)))
    prev_audio_fft = norm_fft.copy()
    is_onset = bool(flux > 0.20)
    
    # 5. Calculate entropy on tonotopic bins for auditory surprise
    prob = tonotopic_bins / (np.sum(tonotopic_bins) + 1e-6)
    entropy = -np.sum(prob * np.log(prob + 1e-6))
    max_entropy = np.log(len(prob)) if len(prob) > 0 else 1.0
    norm_entropy = float(entropy / max_entropy if max_entropy > 0 else 0)
    
    # 6. Spectral Centroid
    freqs = np.linspace(0, fs / 2.0, len(fft_vals))
    spectral_centroid = float(np.sum(freqs * fft_vals) / (np.sum(fft_vals) + 1e-6))
    
    return db, norm_fft.tolist(), norm_entropy, flux, is_onset, f0_pitch, spectral_centroid

def main():
    print(json.dumps({"status": "ready"}), flush=True)
    
    cap_rgb = None
    cap_ir = None
    def init_cameras():
        nonlocal cap_rgb, cap_ir
        try:
            is_windows = sys.platform.startswith('win')
            is_linux = sys.platform.startswith('linux')
            is_macos = sys.platform == 'darwin'

            # Determine optimal video capture backend per OS
            if is_windows:
                backend = cv2.CAP_DSHOW
            elif is_linux:
                backend = cv2.CAP_V4L2
            elif is_macos:
                backend = cv2.CAP_AVFOUNDATION
            else:
                backend = cv2.CAP_ANY

            for dev_idx in [0, 1, 2]:
                if is_linux and not os.path.exists(f'/dev/video{dev_idx}'):
                    continue
                try:
                    c = cv2.VideoCapture(dev_idx, backend) if backend != cv2.CAP_ANY else cv2.VideoCapture(dev_idx)
                    if c is not None and c.isOpened():
                        c.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                        ret, frame = c.read()
                        if ret and frame is not None:
                            if cap_rgb is None:
                                cap_rgb = c
                                print(json.dumps({"status": "info", "message": f"RGB Camera active on device {dev_idx}"}), file=sys.stderr, flush=True)
                            elif cap_ir is None:
                                cap_ir = c
                                print(json.dumps({"status": "info", "message": f"IR Camera active on device {dev_idx}"}), file=sys.stderr, flush=True)
                                break
                        else:
                            c.release()
                    elif c is not None:
                        c.release()
                except Exception:
                    pass
        except Exception:
            pass

    latest_camera_pixels = np.zeros((VISUAL_H, VISUAL_W, 3), dtype=np.float32)
    latest_mean_luma = 0.0
    cam_lock = threading.Lock()

    def camera_worker():
        nonlocal cap_rgb, cap_ir, latest_camera_pixels, latest_mean_luma
        init_cameras()
        rgb_failures = 0
        ir_failures = 0
        reconnect_counter = 0
        while running:
            if cap_rgb is None and cap_ir is None:
                time.sleep(1.0)
                reconnect_counter += 1
                if reconnect_counter >= 5:
                    reconnect_counter = 0
                    init_cameras()
                    rgb_failures = 0
                    ir_failures = 0
                continue

            rgb_frame = None
            ir_frame = None
            if cap_rgb is not None and cap_rgb.isOpened():
                try:
                    cap_rgb.grab()
                    ret1, f1 = cap_rgb.retrieve()
                    if ret1 and f1 is not None:
                        rgb_frame = f1
                        rgb_failures = 0
                    else:
                        rgb_failures += 1
                        if rgb_failures >= 10:
                            cap_rgb.release()
                            cap_rgb = None
                except Exception:
                    rgb_failures += 1
                    if rgb_failures >= 10 and cap_rgb is not None:
                        try: cap_rgb.release()
                        except Exception: pass
                        cap_rgb = None

            if cap_ir is not None and cap_ir.isOpened():
                try:
                    cap_ir.grab()
                    ret2, f2 = cap_ir.retrieve()
                    if ret2 and f2 is not None:
                        ir_frame = f2
                        ir_failures = 0
                    else:
                        ir_failures += 1
                        if ir_failures >= 10:
                            cap_ir.release()
                            cap_ir = None
                except Exception:
                    ir_failures += 1
                    if ir_failures >= 10 and cap_ir is not None:
                        try: cap_ir.release()
                        except Exception: pass
                        cap_ir = None

            active_frame = rgb_frame if rgb_frame is not None else ir_frame
            if active_frame is not None:
                try:
                    small = cv2.resize(active_frame, (VISUAL_W, VISUAL_H))
                    rgb = cv2.cvtColor(small, cv2.COLOR_BGR2RGB)
                    if ir_frame is not None and rgb_frame is not None:
                        ir_small = cv2.cvtColor(cv2.resize(ir_frame, (VISUAL_W, VISUAL_H)), cv2.COLOR_BGR2GRAY)
                        rgb[..., 2] = np.clip(rgb[..., 2] * 0.7 + ir_small * 0.3, 0, 255)
                    pix = (rgb / 255.0).astype(np.float32)
                    luma = float(np.mean(pix))
                    with cam_lock:
                        latest_camera_pixels = pix
                        latest_mean_luma = luma
                except Exception:
                    pass
            time.sleep(0.04)

    cam_thread = threading.Thread(target=camera_worker)
    cam_thread.daemon = True
    cam_thread.start()
    
    audio_queue = collections.deque(maxlen=10)
    def audio_callback(indata, frames, time_info, status):
        audio_queue.append(indata.copy())
        
    stream = None
    active_sample_rate = 48000

    def init_audio():
        nonlocal stream, active_sample_rate
        try:
            # On Linux with PipeWire/Pulse, if default source is a .monitor sink, fix default source
            if sys.platform.startswith('linux'):
                try:
                    import subprocess
                    out = subprocess.check_output(['pactl', 'get-default-source'], text=True, stderr=subprocess.DEVNULL).strip()
                    if '.monitor' in out:
                        sources = subprocess.check_output(['pactl', 'list', 'sources', 'short'], text=True, stderr=subprocess.DEVNULL).splitlines()
                        for s in sources:
                            parts = s.split('\t')
                            if len(parts) >= 2 and '.monitor' not in parts[1]:
                                subprocess.run(['pactl', 'set-default-source', parts[1]], stderr=subprocess.DEVNULL)
                                break
                except Exception:
                    pass

            devices = sd.query_devices()
            candidates = []
            for i, dev in enumerate(devices):
                if dev['max_input_channels'] > 0:
                    name = dev['name']
                    name_lower = name.lower()
                    if 'monitor' in name_lower or 'sink' in name_lower:
                        continue
                    
                    score = 0
                    if 'digital microphone' in name_lower or 'mic1' in name_lower:
                        score = 12
                    elif 'headset' in name_lower or 'bluez_input' in name_lower:
                        score = 10
                    elif 'stereo microphone' in name_lower or 'mic2' in name_lower or 'microphone' in name_lower or 'mic' in name_lower:
                        score = 8
                    elif 'sof-hda-dsp' in name_lower or 'hw:' in name_lower:
                        score = 6
                    elif 'pipewire' in name_lower or 'pulse' in name_lower:
                        score = 4
                    elif 'default' in name_lower:
                        score = 2
                    else:
                        score = 1
                    candidates.append((score, i, name))

            # Add System Default as fallback (score 0)
            candidates.append((0, None, "System Default Capture"))
            candidates.sort(key=lambda x: x[0], reverse=True)

            for score, dev_id, dev_name in candidates:
                for sr in [48000, 44100, 16000]:
                    try:
                        block = int(sr / 24)
                        test_stream = sd.InputStream(device=dev_id, samplerate=sr, channels=1, dtype='float32', blocksize=block, callback=audio_callback)
                        test_stream.start()
                        stream = test_stream
                        active_sample_rate = sr
                        sys.stderr.write(f"[AVDriver] 🎙️ Active Microphone: Device {dev_id} ({dev_name}) @ {sr}Hz\n")
                        return
                    except Exception:
                        stream = None
        except Exception as e:
            sys.stderr.write(f"[AVDriver] Audio stream init warning: {e}\n")

    init_audio()
    
    def stdin_listener():
        global running
        while running:
            line = sys.stdin.readline()
            if not line:
                running = False
                break

    t_stdin = threading.Thread(target=stdin_listener, daemon=True)
    t_stdin.start()
    
    try:
        while running:
            start_time = time.time()
            
            # 1. Read Audio Non-Blockingly
            if audio_queue:
                audio_chunk = audio_queue.popleft()
                db, phonemes, entropy, flux, is_onset, f0_pitch, spectral_centroid = get_audio_features(audio_chunk)
            else:
                db, phonemes, entropy, flux, is_onset, f0_pitch, spectral_centroid = 0.0, [0.0]*128, 0.0, 0.0, False, 0.0, 0.0
                
            # 2. Get latest webcam frame from non-blocking worker thread
            with cam_lock:
                webcam_pixels = latest_camera_pixels.copy()
                mean_luma = latest_mean_luma
                
            webcam_uint8 = (webcam_pixels * 255.0).astype(np.uint8)
            webcam_pixels_b64 = base64.b64encode(webcam_uint8.tobytes()).decode('utf-8')
            
            # Fast JPEG compression for ultra-low IPC and WebSocket payload (drops from 77KB to 3.5KB)
            webcam_bgr = cv2.cvtColor(webcam_uint8, cv2.COLOR_RGB2BGR)
            _, webcam_jpg = cv2.imencode('.jpg', webcam_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 35])
            webcam_jpg_b64 = base64.b64encode(webcam_jpg.tobytes()).decode('utf-8')

            phonemes_list = phonemes.tolist() if isinstance(phonemes, np.ndarray) else phonemes
            # 3. Emit JSON frame payload matching AVBody.js contract
            payload = {
                "type": "av_data",
                "decibels": float(db),
                "phonemes": phonemes_list,
                "entropy": float(entropy),
                "flux": float(flux),
                "is_onset": bool(is_onset),
                "f0_pitch": float(f0_pitch),
                "spectral_centroid": float(spectral_centroid),
                "webcam_pixels_b64": webcam_pixels_b64,
                "webcam_jpg_b64": webcam_jpg_b64,
                "mean_luma": float(mean_luma)
            }
            print(json.dumps(payload), flush=True)
            
            # If live physical mic sound detected (debug only, disabled in production to prevent daemon log bloat)
            # if db > 12.0 and (start_time - getattr(main, 'last_mic_log', 0)) > 2.0:
            #     main.last_mic_log = start_time
            #     pitch_str = f" | F0 Pitch: {f0_pitch:.1f}Hz" if f0_pitch > 0 else ""
            #     sys.stderr.write(f"[AVDriver] 🎙️ Live Mic Signal Ingested: {db:.1f} dB{pitch_str} (Entropy: {entropy:.2f}, Onset: {is_onset})\n")
            
            # Target ~24 FPS
            elapsed = time.time() - start_time
            sleep_time = (1.0 / 24.0) - elapsed
            if sleep_time > 0:
                time.sleep(sleep_time)
                
    except (KeyboardInterrupt, BrokenPipeError, IOError, OSError):
        pass
    except Exception as e:
        try:
            print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
        except:
            pass
    finally:
        if stream is not None:
            try:
                stream.stop()
                stream.close()
            except Exception:
                pass
        if cap_rgb is not None and hasattr(cap_rgb, 'release'):
            cap_rgb.release()
        if cap_ir is not None and hasattr(cap_ir, 'release'):
            cap_ir.release()

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)
