#!/usr/bin/env python3
import sys
import json
import time
import threading
import numpy as np
from PIL import Image

# Mock anki_vector for graceful failure if connection fails or SDK absent
try:
    import anki_vector
    from anki_vector.util import degrees
    HAS_ANKI = True
except ImportError:
    HAS_ANKI = False

VISUAL_DIM = 256
NUM_FFT_BINS = 128
AUDIO_SAMPLE_RATE = 16000

# Global motor intent
motor_intent = {
    'head_angle': 0.0,
    'lift_height': 0.0,
    'left_wheel': 0.0,
    'right_wheel': 0.0
}

def stdin_listener():
    global motor_intent
    while True:
        line = sys.stdin.readline()
        if not line:
            break
        try:
            cmd = json.loads(line)
            if cmd.get('type') == 'motor_vector':
                # Map vector to Anki motors
                # cmd['vector'] = [GazeX, GazeY, MouseX, MouseY, Click, ...]
                vec = cmd.get('vector', [])
                if len(vec) >= 15:
                    # OcularGaze is index 13, 14
                    gaze_y = vec[14]
                    motor_intent['head_angle'] = max(-22, min(45, gaze_y * 45))
                    
                    # Mouse X/Y is index 0, 1
                    fwd = vec[1] * 50.0
                    turn = vec[0] * 25.0
                    motor_intent['left_wheel'] = fwd + turn
                    motor_intent['right_wheel'] = fwd - turn
                    
                    # Left click is index 2
                    click = vec[2]
                    if click > 0.5:
                        motor_intent['lift_height'] = 1.0
                    else:
                        motor_intent['lift_height'] = 0.0
        except:
            pass

def get_audio_features(pcm_data):
    if not pcm_data or len(pcm_data) == 0:
        return 0, [0.0]*NUM_FFT_BINS, 0.0
    
    # Convert to float numpy array
    indata = np.array(pcm_data, dtype=np.float32) / 32768.0
    rms = np.sqrt(np.mean(indata**2))
    db = 20 * np.log10(rms + 1e-6)
    db = float(max(0, min(120, db + 100))) # normalize
    
    fft_vals = np.abs(np.fft.rfft(indata))
    if len(fft_vals) < NUM_FFT_BINS:
        padded = np.zeros(NUM_FFT_BINS)
        padded[:len(fft_vals)] = fft_vals
        fft_vals = padded
        
    raw_fft_bins = fft_vals[:NUM_FFT_BINS]
    log_fft = np.log10(raw_fft_bins + 1e-6)
    range_val = np.max(log_fft) - np.min(log_fft)
    if range_val > 0:
        log_fft = (log_fft - np.min(log_fft)) / range_val
    else:
        log_fft = np.zeros(NUM_FFT_BINS)
        
    prob = fft_vals / (np.sum(fft_vals) + 1e-6)
    entropy = -np.sum(prob * np.log(prob + 1e-6))
    max_entropy = np.log(len(prob))
    norm_entropy = float(entropy / max_entropy if max_entropy > 0 else 0)
    
    return db, log_fft.tolist(), norm_entropy

def main():
    if not HAS_ANKI:
        print(json.dumps({"status": "error", "message": "anki_vector module not found."}), flush=True)
        sys.exit(1)
        
    robot = None
    try:
        # We start the robot connection
        # Requires ANKI_ROBOT_IP, ANKI_ROBOT_NAME, ANKI_ROBOT_CERT, etc. to be set in env or ~/.anki_vector
        robot = anki_vector.Robot(enable_camera_feed=True, enable_audio_feed=True)
        robot.connect()
    except Exception as e:
        print(json.dumps({"status": "error", "message": f"Failed to connect to Vector: {str(e)}"}), flush=True)
        sys.exit(1)
        
    print(json.dumps({"status": "ready"}), flush=True)
    
    t = threading.Thread(target=stdin_listener, daemon=True)
    t.start()
    
    try:
        while True:
            start_time = time.time()
            
            # 1. Apply Motors
            robot.motors.set_wheel_motors(motor_intent['left_wheel'], motor_intent['right_wheel'])
            robot.behavior.set_head_angle(degrees(motor_intent['head_angle']))
            robot.behavior.set_lift_height(motor_intent['lift_height'])
            
            # 2. Get Audio Chunk
            pcm_data = []
            if robot.audio.stream_enabled and robot.audio.has_audio_data:
                # the sdk provides raw 16-bit PCM chunks
                # we just grab whatever is in the buffer
                try:
                    audio_chunk = robot.audio.latest_audio_chunk
                    if audio_chunk:
                        # Convert bytes to int16 list
                        pcm_data = np.frombuffer(audio_chunk, dtype=np.int16).tolist()
                except:
                    pass
                    
            db, fft_bins, entropy = get_audio_features(pcm_data)
            
            # 3. Get Vision (Compressed Base64 JPEG buffer to eliminate IPC serialization bottlenecks)
            img_b64 = None
            if robot.camera.latest_image:
                try:
                    import io, base64
                    img = robot.camera.latest_image.raw_image
                    img_resized = img.resize((VISUAL_DIM, VISUAL_DIM))
                    buf = io.BytesIO()
                    img_resized.save(buf, format='JPEG', quality=75)
                    img_b64 = base64.b64encode(buf.getvalue()).decode('ascii')
                except:
                    pass
                    
            msg = {
                "type": "vector_av",
                "image_b64": img_b64,
                "decibels": db,
                "fft_bins": fft_bins,
                "entropy": entropy
            }
            
            print(json.dumps(msg), flush=True)
            
            # 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:
        pass
    finally:
        if robot is not None:
            try:
                robot.disconnect()
            except Exception:
                pass

if __name__ == '__main__':
    main()
