#!/usr/bin/env python3
"""
Stage 3 + Faces – Cozmo + HoboStreamer
Procedural eyes by default + simple static face cycling

Public version – replace STREAM_KEY with your own.
"""

import time
import threading
import json
import os
from queue import Queue, Empty
from pathlib import Path

import pygame
import cv2
import numpy as np
from PIL import Image
import pycozmo
from pycozmo import protocol_encoder

try:
    import websocket
except ImportError:
    print("Missing websocket-client. Run: pip install websocket-client")
    raise

# ========================== TUNABLES ==========================
DRIVE_SPEED_FAST   = 140
DRIVE_SPEED_MEDIUM = 88
DRIVE_SPEED_SLOW   = 55

JUMPY_BURST_DURATION = 0.07
JUMPY_PAUSE_TIME     = 0.03

HEAD_STEP            = 0.10
LIFT_STEP            = 6
STICK_DEADZONE       = 0.08

CAMERA_GAIN          = 3.90
CAMERA_EXPOSURE      = 100

RECONNECT_DELAY      = 4.0
WS_RECONNECT_DELAY   = 5.0

# Face folder – drop 128x32 BMP or PNG files here
FACES_DIR = Path(r"C:\Users\matth\Desktop\cozmo_projects\faces")

# ←←← PUT YOUR HOBOSTREAMER STREAM KEY HERE ←←←
STREAM_KEY = "YOUR_STREAM_KEY_HERE"
WS_URL = f"wss://hobostreamer.com/ws/control?mode=hardware&stream_key={STREAM_KEY}"
# ==============================================================

GEARS = [
    ("FAST",   DRIVE_SPEED_FAST),
    ("MEDIUM", DRIVE_SPEED_MEDIUM),
    ("SLOW",   DRIVE_SPEED_SLOW),
]

latest_image = None
MIRROR_MODE = True
VIEWER_ENABLED = True
DRIVING_MODE = "jumpy"
gear_index = 1
ir_light_on = False

# Face state
procedural_on = True
static_faces = []
static_names = []
current_face_idx = -1

remote_queue = Queue(maxsize=40)

throttle_min = None
throttle_max = None
throttle_calibrated = False


def on_camera(cli, image):
    global latest_image
    latest_image = image.copy()


def clamp(v, lo, hi):
    return max(lo, min(hi, v))


def apply_camera_settings(cli):
    try:
        pkt = protocol_encoder.SetCameraParams(
            gain=float(CAMERA_GAIN),
            exposure_ms=int(CAMERA_EXPOSURE),
            auto_exposure_enabled=False
        )
        cli.conn.send(pkt)
    except Exception as e:
        print(f"Camera settings error: {e}")


def set_ir(cli, on: bool):
    global ir_light_on
    ir_light_on = on
    try:
        cli.set_head_light(enable=on)
        print(f"IR light: {'ON' if on else 'OFF'}")
    except Exception as e:
        print(f"IR error: {e}")


def load_static_faces():
    global static_faces, static_names
    static_faces = []
    static_names = []

    if not FACES_DIR.exists():
        FACES_DIR.mkdir(parents=True, exist_ok=True)
        print(f"Created faces folder: {FACES_DIR}")
        print("Drop 128x32 BMP or PNG files in there.")
        return

    files = sorted(
        list(FACES_DIR.glob("*.bmp")) +
        list(FACES_DIR.glob("*.png")) +
        list(FACES_DIR.glob("*.BMP")) +
        list(FACES_DIR.glob("*.PNG"))
    )

    for path in files:
        try:
            im = Image.open(path).convert("1").resize((128, 32), Image.NEAREST)
            static_faces.append(im)
            static_names.append(path.name)
        except Exception as e:
            print(f"Could not load face {path.name}: {e}")

    print(f"Loaded {len(static_faces)} static face(s)")


def show_procedural(cli):
    global procedural_on, current_face_idx
    try:
        cli.enable_procedural_face(True)
        procedural_on = True
        current_face_idx = -1
        print("Procedural eyes ON")
    except Exception as e:
        print(f"Procedural face error: {e}")


def show_static_face(cli, idx):
    global procedural_on, current_face_idx
    if not static_faces:
        print("No static faces loaded")
        return
    idx = idx % len(static_faces)
    try:
        cli.enable_procedural_face(False)
        cli.display_image(static_faces[idx])
        procedural_on = False
        current_face_idx = idx
        print(f"Face → {static_names[idx]}")
    except Exception as e:
        print(f"Display face error: {e}")


def next_static_face(cli):
    if not static_faces:
        print("No static faces loaded")
        return
    next_idx = (current_face_idx + 1) % len(static_faces)
    show_static_face(cli, next_idx)


def on_ws_message(ws, raw):
    try:
        msg = json.loads(raw)
        t = msg.get("type")
        if t in ("command", "key_down", "key_up", "video_click"):
            remote_queue.put(msg)
    except Exception as e:
        print(f"WS parse error: {e}")


def ws_thread_fn():
    while True:
        try:
            print("Connecting to HoboStreamer...")
            ws = websocket.WebSocketApp(
                WS_URL,
                on_message=on_ws_message,
                on_open=lambda ws: print("HoboStreamer connected"),
                on_close=lambda ws, *a: print("HoboStreamer disconnected"),
                on_error=lambda ws, e: print(f"WS error: {e}"),
            )
            ws.run_forever(ping_interval=25, ping_timeout=10)
        except Exception as e:
            print(f"WS thread: {e}")
        time.sleep(WS_RECONNECT_DELAY)


def handle_remote(cli, msg, head, lift, drive_speed):
    global DRIVING_MODE, gear_index, VIEWER_ENABLED

    if not VIEWER_ENABLED:
        return head, lift

    cmd = (msg.get("command") or "").lower().strip()
    if not cmd:
        return head, lift

    print(f"Viewer → {cmd}")

    if cmd in ("forward", "backward", "turn_left", "turn_right"):
        if cmd == "forward":
            l = r = drive_speed
        elif cmd == "backward":
            l = r = -drive_speed
        elif cmd == "turn_left":
            l, r = -int(drive_speed * 0.75), int(drive_speed * 0.75)
        else:
            l, r = int(drive_speed * 0.75), -int(drive_speed * 0.75)

        if DRIVING_MODE == "smooth":
            cli.drive_wheels(l, r, duration=0.35)
        else:
            cli.drive_wheels(l, r, duration=JUMPY_BURST_DURATION)
            time.sleep(JUMPY_PAUSE_TIME)

    elif cmd == "head_up":
        head = clamp(head + HEAD_STEP, pycozmo.MIN_HEAD_ANGLE.radians, pycozmo.MAX_HEAD_ANGLE.radians)
        cli.set_head_angle(head)
        handle_remote._remote_head_until = time.time() + 1.2

    elif cmd == "head_down":
        head = clamp(head - HEAD_STEP, pycozmo.MIN_HEAD_ANGLE.radians, pycozmo.MAX_HEAD_ANGLE.radians)
        cli.set_head_angle(head)
        handle_remote._remote_head_until = time.time() + 1.2

    elif cmd == "lift_up":
        lift = clamp(lift + 20, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
        cli.set_lift_height(lift)

    elif cmd == "lift_down":
        lift = clamp(lift - 20, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
        cli.set_lift_height(lift)

    elif cmd in ("ir_light", "l"):
        set_ir(cli, not ir_light_on)

    elif cmd in ("cycle_gear", "z"):
        gear_index = (gear_index + 1) % len(GEARS)
        print(f"Gear → {GEARS[gear_index][0]}")

    elif cmd in ("toggle_mode", "x"):
        DRIVING_MODE = "smooth" if DRIVING_MODE == "jumpy" else "jumpy"
        print(f"Mode → {DRIVING_MODE.upper()}")

    elif cmd in ("next_face", "face", "cycle_face"):
        next_static_face(cli)

    elif cmd in ("procedural", "procedural_eyes", "default_eyes", "eyes"):
        show_procedural(cli)

    elif cmd.startswith("face_"):
        try:
            idx = int(cmd.split("_")[1])
            show_static_face(cli, idx)
        except Exception:
            pass

    elif cmd in ("stop", "emergency_stop", "space"):
        cli.drive_wheels(0, 0)

    return head, lift


def run_once(joy):
    global CAMERA_GAIN, CAMERA_EXPOSURE, MIRROR_MODE, VIEWER_ENABLED
    global DRIVING_MODE, gear_index, throttle_min, throttle_max, throttle_calibrated
    global procedural_on

    print("Connecting to Cozmo...")
    with pycozmo.connect(enable_procedural_face=True) as cli:
        cli.enable_camera(enable=True, color=True)
        cli.add_handler(pycozmo.event.EvtNewRawCameraImage, on_camera)
        time.sleep(0.7)
        apply_camera_settings(cli)
        set_ir(cli, False)
        show_procedural(cli)

        print("Cozmo ready")
        print(f"  Gear: {GEARS[gear_index][0]}   Mode: {DRIVING_MODE}   Viewer: {VIEWER_ENABLED}")

        head = 0.0
        lift = pycozmo.MIN_LIFT_HEIGHT.mm
        clock = pygame.time.Clock()
        running = True
        last_battery = 0
        battery_v = 0.0
        battery_pct = 0

        if not hasattr(handle_remote, "_remote_head_until"):
            handle_remote._remote_head_until = 0

        while running:
            now = time.time()
            drive_speed = GEARS[gear_index][1]

            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    running = False
                elif e.type == pygame.KEYDOWN:
                    if e.key == pygame.K_ESCAPE:
                        running = False
                    elif e.key in (pygame.K_EQUALS, pygame.K_PLUS):
                        CAMERA_GAIN = clamp(CAMERA_GAIN + 0.15, 0.1, 4.5)
                        apply_camera_settings(cli)
                        print(f"Gain → {CAMERA_GAIN:.2f}")
                    elif e.key == pygame.K_MINUS:
                        CAMERA_GAIN = clamp(CAMERA_GAIN - 0.15, 0.1, 4.5)
                        apply_camera_settings(cli)
                        print(f"Gain → {CAMERA_GAIN:.2f}")
                    elif e.key == pygame.K_RIGHTBRACKET:
                        CAMERA_EXPOSURE = clamp(CAMERA_EXPOSURE + 5, 1, 120)
                        apply_camera_settings(cli)
                        print(f"Exposure → {CAMERA_EXPOSURE} ms")
                    elif e.key == pygame.K_LEFTBRACKET:
                        CAMERA_EXPOSURE = clamp(CAMERA_EXPOSURE - 5, 1, 120)
                        apply_camera_settings(cli)
                        print(f"Exposure → {CAMERA_EXPOSURE} ms")
                    elif e.key == pygame.K_m:
                        MIRROR_MODE = not MIRROR_MODE
                        print(f"Mirror: {MIRROR_MODE}")
                    elif e.key == pygame.K_c:
                        VIEWER_ENABLED = not VIEWER_ENABLED
                        print(f"Viewer control: {'ON' if VIEWER_ENABLED else 'OFF'}")
                    elif e.key == pygame.K_x:
                        DRIVING_MODE = "smooth" if DRIVING_MODE == "jumpy" else "jumpy"
                        print(f"Mode → {DRIVING_MODE.upper()}")
                    elif e.key == pygame.K_z:
                        gear_index = (gear_index + 1) % len(GEARS)
                        print(f"Gear → {GEARS[gear_index][0]}")
                    elif e.key == pygame.K_l:
                        set_ir(cli, not ir_light_on)
                    elif e.key == pygame.K_n:
                        next_static_face(cli)
                    elif e.key == pygame.K_o:
                        show_procedural(cli)
                    elif e.key == pygame.K_p:
                        load_static_faces()

            pygame.event.pump()
            keys = pygame.key.get_pressed()

            while True:
                try:
                    msg = remote_queue.get_nowait()
                    head, lift = handle_remote(cli, msg, head, lift, drive_speed)
                except Empty:
                    break

            left = right = 0.0
            if joy:
                y = joy.get_axis(1)
                twist = joy.get_axis(2)
                if abs(y) < STICK_DEADZONE: y = 0.0
                if abs(twist) < STICK_DEADZONE: twist = 0.0
                forward = -y * drive_speed
                turn = -twist * drive_speed
                if MIRROR_MODE:
                    turn = -turn
                left = forward - turn
                right = forward + turn

            if keys[pygame.K_w] or keys[pygame.K_UP]:
                left = right = drive_speed
            elif keys[pygame.K_s] or keys[pygame.K_DOWN]:
                left = right = -drive_speed
            if keys[pygame.K_a] or keys[pygame.K_LEFT]:
                left, right = -drive_speed, drive_speed
            elif keys[pygame.K_d] or keys[pygame.K_RIGHT]:
                left, right = drive_speed, -drive_speed

            if keys[pygame.K_SPACE]:
                left = right = 0

            if abs(left) > 5 or abs(right) > 5:
                if DRIVING_MODE == "smooth":
                    cli.drive_wheels(left, right)
                else:
                    cli.drive_wheels(left, right, duration=JUMPY_BURST_DURATION)
                    time.sleep(JUMPY_PAUSE_TIME)
            else:
                cli.drive_wheels(0, 0)

            remote_head_priority = getattr(handle_remote, "_remote_head_until", 0)

            if keys[pygame.K_q]:
                head = clamp(head + 0.06, pycozmo.MIN_HEAD_ANGLE.radians, pycozmo.MAX_HEAD_ANGLE.radians)
            if keys[pygame.K_e]:
                head = clamp(head - 0.06, pycozmo.MIN_HEAD_ANGLE.radians, pycozmo.MAX_HEAD_ANGLE.radians)

            if joy and joy.get_numaxes() >= 4 and now > remote_head_priority:
                raw = joy.get_axis(3)
                if throttle_min is None:
                    throttle_min = throttle_max = raw
                else:
                    throttle_min = min(throttle_min, raw)
                    throttle_max = max(throttle_max, raw)
                if (throttle_max - throttle_min) > 0.25:
                    norm = (raw - throttle_min) / (throttle_max - throttle_min)
                    head = pycozmo.MIN_HEAD_ANGLE.radians + norm * (
                        pycozmo.MAX_HEAD_ANGLE.radians - pycozmo.MIN_HEAD_ANGLE.radians)
                    if not throttle_calibrated:
                        throttle_calibrated = True
                        print("Throttle calibrated")

            cli.set_head_angle(head)

            if keys[pygame.K_r]:
                lift = clamp(lift + LIFT_STEP, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
                cli.set_lift_height(lift)
            if keys[pygame.K_f]:
                lift = clamp(lift - LIFT_STEP, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
                cli.set_lift_height(lift)

            if joy:
                hat = joy.get_hat(0)
                if hat[1] == 1:
                    lift = clamp(lift + LIFT_STEP, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
                    cli.set_lift_height(lift)
                elif hat[1] == -1:
                    lift = clamp(lift - LIFT_STEP, pycozmo.MIN_LIFT_HEIGHT.mm, pycozmo.MAX_LIFT_HEIGHT.mm)
                    cli.set_lift_height(lift)

            if now - last_battery > 4.0:
                try:
                    battery_v = getattr(cli, "battery_voltage", 0.0) or 0.0
                    battery_pct = int(clamp((battery_v - 3.0) * 100 / 1.2, 0, 100))
                except Exception:
                    battery_v = 0.0
                    battery_pct = 0
                last_battery = now

            if latest_image is not None:
                img = cv2.cvtColor(np.array(latest_image), cv2.COLOR_RGB2BGR)
                cv2.imshow("Cozmo Eye", img)

            hud = np.zeros((180, 320, 3), dtype=np.uint8)
            cv2.putText(hud, f"Batt {battery_pct}%  {battery_v:.2f}V", (10, 28),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0) if battery_pct > 30 else (0, 100, 255), 2)
            cv2.putText(hud, f"Gear: {GEARS[gear_index][0]}", (10, 55),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 100), 2)
            cv2.putText(hud, f"Mode: {DRIVING_MODE.upper()}", (10, 82),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 255), 2)
            cv2.putText(hud, f"Viewer: {'ON' if VIEWER_ENABLED else 'OFF'}", (10, 109),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0) if VIEWER_ENABLED else (0, 0, 255), 2)
            face_txt = "Eyes: Procedural" if procedural_on else f"Face: {static_names[current_face_idx] if current_face_idx >= 0 else '?'}"
            cv2.putText(hud, face_txt, (10, 136),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.50, (200, 200, 255), 1)
            cv2.putText(hud, f"IR: {'ON' if ir_light_on else 'OFF'}", (10, 163),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.50, (200, 200, 200), 1)
            cv2.imshow("HUD", hud)

            cv2.waitKey(1)
            clock.tick(35)

        cli.drive_wheels(0, 0)
        cli.enable_camera(enable=False)
        set_ir(cli, False)
        print("Session ended cleanly.")


def main():
    load_static_faces()

    pygame.init()
    pygame.joystick.init()
    pygame.display.set_mode((320, 180))
    pygame.display.set_caption("Stage 3 Cozmo + Faces – ESC to quit")

    joy = None
    if pygame.joystick.get_count() > 0:
        joy = pygame.joystick.Joystick(0)
        joy.init()
        print(f"Joystick: {joy.get_name()}")
    else:
        print("No joystick – keyboard only")

    cv2.namedWindow("Cozmo Eye", cv2.WINDOW_NORMAL)
    cv2.namedWindow("HUD", cv2.WINDOW_NORMAL)

    t = threading.Thread(target=ws_thread_fn, daemon=True)
    t.start()

    print("\nControls:")
    print("  WASD/Arrows   Drive")
    print("  Q/E           Head")
    print("  Throttle      Head")
    print("  R/F + Hat     Lift")
    print("  Z / X         Gear / Mode")
    print("  C             Viewer on/off")
    print("  L             IR light")
    print("  N             Next static face")
    print("  O             Procedural eyes")
    print("  P             Reload faces folder")
    print("  M             Mirror")
    print("  Space         Stop")
    print("  ESC           Quit\n")

    while True:
        try:
            run_once(joy)
            break
        except Exception as e:
            print(f"Lost connection: {e}")
            print(f"Reconnecting in {RECONNECT_DELAY}s...")
            time.sleep(RECONNECT_DELAY)

    cv2.destroyAllWindows()
    pygame.quit()


if __name__ == "__main__":
    main()