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

if hasattr(signal, 'SIGPIPE'):
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)

import json
import time
import threading
import math
from pynput.mouse import Controller as MouseController, Listener as MouseListener, Button
from pynput.keyboard import Controller as KeyboardController, Key, KeyCode, Listener as KeyboardListener

mouse = MouseController()
keyboard = KeyboardController()

# A map of string modifiers to pynput Keys
MOD_MAP = {
    'KEY_LEFTMETA': Key.cmd,
    'KEY_RIGHTMETA': Key.cmd_r,
    'KEY_LEFTSHIFT': Key.shift,
    'KEY_RIGHTSHIFT': Key.shift_r,
    'KEY_LEFTCTRL': Key.ctrl,
    'KEY_RIGHTCTRL': Key.ctrl_r,
    'KEY_LEFTALT': Key.alt,
    'KEY_RIGHTALT': Key.alt_r,
}

KEY_MAP = {
    'KEY_ENTER': Key.enter,
    'KEY_SPACE': Key.space,
    'KEY_TAB': Key.tab,
    'KEY_BACKSPACE': Key.backspace,
    'KEY_ESC': Key.esc,
    'KEY_UP': Key.up,
    'KEY_DOWN': Key.down,
    'KEY_LEFT': Key.left,
    'KEY_RIGHT': Key.right,
    'KEY_DELETE': Key.delete,
    'KEY_HOME': Key.home,
    'KEY_END': Key.end,
    'KEY_PAGEUP': Key.page_up,
    'KEY_PAGEDOWN': Key.page_down,
    'KEY_INSERT': Key.insert,
    'KEY_CAPSLOCK': Key.caps_lock,
    'KEY_LEFTSHIFT': Key.shift,
    'KEY_RIGHTSHIFT': Key.shift_r,
    'KEY_LEFTCTRL': Key.ctrl,
    'KEY_RIGHTCTRL': Key.ctrl_r,
    'KEY_LEFTALT': Key.alt,
    'KEY_RIGHTALT': Key.alt_r,
    'KEY_LEFTMETA': Key.cmd,
    'KEY_RIGHTMETA': Key.cmd_r,
    'KEY_NUMLOCK': getattr(Key, 'num_lock', Key.esc),
    'KEY_SCROLLLOCK': getattr(Key, 'scroll_lock', Key.esc),
    'KEY_PAUSE': getattr(Key, 'pause', Key.esc),
    'KEY_SYSRQ': getattr(Key, 'print_screen', Key.esc),
    'KEY_COMPOSE': getattr(Key, 'menu', Key.esc),
    # Symbols & Punctuation
    'KEY_DOT': KeyCode.from_char('.'),
    'KEY_COMMA': KeyCode.from_char(','),
    'KEY_SEMICOLON': KeyCode.from_char(';'),
    'KEY_COLON': KeyCode.from_char(':'),
    'KEY_APOSTROPHE': KeyCode.from_char("'"),
    'KEY_DOUBLEQUOTE': KeyCode.from_char('"'),
    'KEY_EXCLAMATION': KeyCode.from_char('!'),
    'KEY_QUESTION': KeyCode.from_char('?'),
    'KEY_GRAVE': KeyCode.from_char('`'),
    'KEY_MINUS': KeyCode.from_char('-'),
    'KEY_EQUAL': KeyCode.from_char('='),
    'KEY_SLASH': KeyCode.from_char('/'),
    'KEY_BACKSLASH': KeyCode.from_char('\\'),
    'KEY_LEFTBRACE': KeyCode.from_char('['),
    'KEY_RIGHTBRACE': KeyCode.from_char(']'),
    'KEY_KPDOT': KeyCode.from_char('.'),
    'KEY_KPENTER': Key.enter,
    'KEY_KPPLUS': KeyCode.from_char('+'),
    'KEY_KPMINUS': KeyCode.from_char('-'),
    'KEY_KPASTERISK': KeyCode.from_char('*'),
    'KEY_KPSLASH': KeyCode.from_char('/'),
}

# --- Reflex Arc State ---
state_lock = threading.Lock()
target_x = None
target_y = None
adrenaline = 0.0
serotonin = 0.0
is_saccading = False
last_sdo_actuation_time = 0.0
last_sdo_keystroke_time = 0.0
last_sdo_keystroke_name = ""
last_human_move_emit = 0.0
human_override_until = 0.0

def on_move(x, y):
    global last_human_move_emit, human_override_until, target_x, target_y
    now = time.time()
    if now - last_sdo_actuation_time > 0.06:
        human_override_until = now + 2.5
        with state_lock:
            target_x = None
            target_y = None
        # Rate-limit human move emissions to ~40Hz
        if now - last_human_move_emit > 0.025:
            last_human_move_emit = now
            print(json.dumps({"type": "position", "x": x, "y": y, "origin": "HUMAN"}), flush=True)

def on_click(x, y, button, pressed):
    global human_override_until, target_x, target_y
    now = time.time()
    if now - last_sdo_actuation_time < 0.15:
        return
    human_override_until = now + 2.5
    with state_lock:
        target_x = None
        target_y = None
    if pressed:
        btn_name = "LEFT" if button == Button.left else ("RIGHT" if button == Button.right else "MIDDLE")
        print(json.dumps({"type": "click", "button": btn_name, "x": x, "y": y, "origin": "HUMAN"}), flush=True)

def on_scroll(x, y, dx, dy):
    global human_override_until, target_x, target_y
    now = time.time()
    if now - last_sdo_actuation_time < 0.15:
        return
    human_override_until = now + 2.5
    with state_lock:
        target_x = None
        target_y = None
    print(json.dumps({"type": "position", "x": x, "y": y, "origin": "HUMAN"}), flush=True)

def on_press(key):
    global human_override_until, target_x, target_y, last_sdo_keystroke_time, last_sdo_keystroke_name
    now = time.time()
    try:
        if hasattr(key, 'char') and key.char:
            char_val = key.char
            key_name = f"KEY_{char_val.upper()}"
        else:
            key_name = str(key).replace("Key.", "KEY_").upper()
            char_val = ""
            
        # Reafference filter: SDO's own commanded keystroke within 250ms is NOT a human teacher demonstration!
        if (now - last_sdo_keystroke_time < 0.25) and (key_name == last_sdo_keystroke_name):
            last_sdo_keystroke_time = 0.0
            print(json.dumps({"type": "keystroke", "key": key_name, "char": char_val, "origin": "SDO"}), flush=True)
            return

        human_override_until = now + 2.5
        with state_lock:
            target_x = None
            target_y = None
        print(json.dumps({"type": "keystroke", "key": key_name, "char": char_val, "origin": "HUMAN"}), flush=True)
    except Exception:
        pass

def parse_key(key_str):
    if not key_str:
        return None
    if key_str in KEY_MAP:
        return KEY_MAP[key_str]
    # Direct single character (e.g. 'a', '1', ' ')
    if len(key_str) == 1:
        return KeyCode.from_char(key_str)
    # KEY_A - KEY_Z or KEY_0 - KEY_9
    if key_str.startswith('KEY_') and len(key_str) == 5:
        char = key_str[4].lower()
        return KeyCode.from_char(char)
    # KEY_KP0 - KEY_KP9
    if key_str.startswith('KEY_KP') and len(key_str) == 7 and key_str[6].isdigit():
        return KeyCode.from_char(key_str[6])
    # Function keys KEY_F1 - KEY_F12
    if key_str.startswith('KEY_F') and len(key_str) in (6, 7):
        f_num = key_str[5:]
        if f_num.isdigit() and hasattr(Key, f'f{f_num}'):
            return getattr(Key, f'f{f_num}')
    # Generic strip prefix fallback
    clean = key_str.replace('KEY_', '').lower()
    if len(clean) == 1:
        return KeyCode.from_char(clean)
    return None

def reflex_arc_loop():
    global target_x, target_y, is_saccading, adrenaline, serotonin, human_override_until
    vel_x = 0.0
    vel_y = 0.0
    
    while True:
        if time.time() < human_override_until:
            time.sleep(1.0 / 60.0)
            continue

        with state_lock:
            tx, ty = target_x, target_y
            adr, ser = adrenaline, serotonin
            
        if tx is not None and ty is not None:
            curr_pos = mouse.position
            curr_x, curr_y = curr_pos[0], curr_pos[1]
            
            error_x = tx - curr_x
            error_y = ty - curr_y
            dist = math.hypot(error_x, error_y)
            
            if dist > 300: # Saccade threshold
                with state_lock:
                    is_saccading = True
                
                last_sdo_actuation_time = time.time()
                mouse.position = (tx, ty)
                print(json.dumps({"type": "position", "x": tx, "y": ty, "origin": "SDO", "saccade": True}), flush=True)
                
                with state_lock:
                    is_saccading = False
                vel_x = 0.0
                vel_y = 0.0
            else:
                # PID / Spring-Damper
                # Adrenaline increases spring stiffness (P)
                k_p = 0.08 + (adr * 0.25)
                
                # Serotonin increases damping (D)
                k_d = 0.25 + (ser * 0.45)
                
                acc_x = (error_x * k_p) - (vel_x * k_d)
                acc_y = (error_y * k_p) - (vel_y * k_d)
                
                vel_x += acc_x
                vel_y += acc_y
                
                # Sub-pixel accumulator for smooth motion without truncation
                move_dx = int(round(vel_x)) if abs(vel_x) >= 1.0 else (1 if vel_x > 0.15 else (-1 if vel_x < -0.15 else 0))
                move_dy = int(round(vel_y)) if abs(vel_y) >= 1.0 else (1 if vel_y > 0.15 else (-1 if vel_y < -0.15 else 0))
                
                if move_dx != 0 or move_dy != 0:
                    last_sdo_actuation_time = time.time()
                    mouse.move(move_dx, move_dy)
                    new_pos = mouse.position
                    print(json.dumps({"type": "position", "x": new_pos[0], "y": new_pos[1], "origin": "SDO", "saccade": False}), flush=True)
        
        time.sleep(1.0 / 60.0)

accum_dx = 0.0
accum_dy = 0.0

def main():
    global accum_dx, accum_dy, last_sdo_actuation_time, human_override_until
    print('{"status": "ready"}', flush=True)

    mouse_listener = MouseListener(on_move=on_move, on_click=on_click, on_scroll=on_scroll)
    mouse_listener.start()
    
    keyboard_listener = KeyboardListener(on_press=on_press)
    keyboard_listener.start()
    
    reflex_thread = threading.Thread(target=reflex_arc_loop, daemon=True)
    reflex_thread.start()

    try:
        while True:
            line = sys.stdin.readline()
            if not line:
                break
            
            try:
                cmd = json.loads(line)
                action = cmd.get('type')
                
                if time.time() < human_override_until:
                    continue

                if action == 'expect':
                    global target_x, target_y, adrenaline, serotonin
                    with state_lock:
                        target_x = cmd.get('target_x')
                        target_y = cmd.get('target_y')
                        adrenaline = cmd.get('adrenaline', 0.0)
                        serotonin = cmd.get('serotonin', 0.0)
                
                elif action == 'move':
                    with state_lock:
                        target_x = None
                        target_y = None
                    accum_dx += float(cmd.get('dx', 0.0))
                    accum_dy += float(cmd.get('dy', 0.0))
                    step_x = int(accum_dx)
                    step_y = int(accum_dy)
                    if step_x != 0 or step_y != 0:
                        last_sdo_actuation_time = time.time()
                        mouse.move(step_x, step_y)
                        accum_dx -= step_x
                        accum_dy -= step_y
                
                elif action == 'click':
                    last_sdo_actuation_time = time.time()
                    btn = cmd.get('button', 'left')
                    if btn == 'right':
                        mouse.click(Button.right, 1)
                    else:
                        mouse.click(Button.left, 1)
                
                elif action == 'scroll':
                    last_sdo_actuation_time = time.time()
                    mouse.scroll(cmd.get('dx', 0), cmd.get('dy', 0))
                
                elif action == 'combo':
                    mods = cmd.get('mods', [])
                    key_str = cmd.get('key')
                    
                    pressed_mods = []
                    for mod in mods:
                        if mod in MOD_MAP:
                            k = MOD_MAP[mod]
                            keyboard.press(k)
                            pressed_mods.append(k)
                    
                    if key_str:
                        k = parse_key(key_str)
                        if k:
                            last_sdo_keystroke_time = time.time()
                            last_sdo_keystroke_name = key_str
                            keyboard.press(k)
                            keyboard.release(k)
                    
                    for k in reversed(pressed_mods):
                        keyboard.release(k)
                
                elif action == 'type':
                    text = cmd.get('text', '')
                    keyboard.type(text)

            except Exception as e:
                print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
                
    except KeyboardInterrupt:
        pass
    finally:
        mouse_listener.stop()
        keyboard_listener.stop()

if __name__ == '__main__':
    main()
