"""
spot_bridge.py - command bridge with safety layer.

Receives gesture command events over UDP from predict_headless.py,
applies safety checks, and forwards accepted commands to a robot backend.

Backends:
  FakeSpot  - prints actions (default, safe for testing)
  RealSpot  - commands the real robot via the Boston Dynamics SDK

To use the real robot:
  export SPOT_BACKEND=real
  export SPOT_IP=192.168.50.3
  export SPOT_USER=<username>
  export SPOT_PASSWORD=<password>
  python spot_bridge.py
"""

import os
import socket
import json
import time

# ---------- Safety configuration ----------
WHITELIST = {"STOP", "SIT", "STAND", "WALK_10M", "TURN_180", "FORWARD_5M"}
MIN_CMD_GAP = 2.0          # seconds between accepted commands (STOP exempt)
HEARTBEAT_TIMEOUT = 2.0    # no heartbeat this long -> auto-STOP
LISTEN_ADDR = ("127.0.0.1", 5005)


# ---------- Robot backends ----------
class FakeSpot:
    """Stands in for the real robot. Same interface as RealSpot."""

    def stop(self):     print(">>> SPOT: STOP (halt all motion)")
    def sit(self):      print(">>> SPOT: SIT")
    def stand(self):    print(">>> SPOT: STAND")
    def walk_10m(self): print(">>> SPOT: WALK FORWARD 10 METRES")
    def close(self):    pass
    def turn_180(self):   print(">>> SPOT: TURN 180 degrees")
    def forward_5m(self): print(">>> SPOT: WALK FORWARD 5 m")

class RealSpot:
    """Commands the real robot via the Boston Dynamics SDK."""

    def __init__(self, hostname, username, password):
        import bosdyn.client
        from bosdyn.client.lease import LeaseClient, LeaseKeepAlive
        from bosdyn.client.robot_command import RobotCommandClient
        from bosdyn.client.robot_state import RobotStateClient

        sdk = bosdyn.client.create_standard_sdk("PiGestureController")
        self.robot = sdk.create_robot(hostname)
        self.robot.authenticate(username, password)
        self.robot.time_sync.wait_for_sync()

        if self.robot.is_estopped():
            raise RuntimeError(
                "Spot is E-Stopped. Start an external E-Stop endpoint "
                "(estop_nogui.py) before commanding motors."
            )

        self.command_client = self.robot.ensure_client(
            RobotCommandClient.default_service_name)
        self.state_client = self.robot.ensure_client(
            RobotStateClient.default_service_name)
        self.lease_client = self.robot.ensure_client(
            LeaseClient.default_service_name)
        self.lease_client.take()
        self.lease_keepalive = LeaseKeepAlive(
            self.lease_client, must_acquire=False, return_at_exit=True)

        print("RealSpot: connected, time synced, lease acquired")

    def power_on(self):
        self.robot.power_on(timeout_sec=20)
        if not self.robot.is_powered_on():
            raise RuntimeError("Spot failed to power on.")

    def stand(self):
        from bosdyn.client.robot_command import blocking_stand
        self.power_on()
        blocking_stand(self.command_client, timeout_sec=10)

    def sit(self):
        from bosdyn.client.robot_command import RobotCommandBuilder
        self.command_client.robot_command(
            RobotCommandBuilder.synchro_sit_command())

    def stop(self):
        from bosdyn.client.robot_command import RobotCommandBuilder
        self.command_client.robot_command(
            RobotCommandBuilder.stop_command())

    def walk_10m(self, distance_m=10.0):
        from bosdyn.client.robot_command import RobotCommandBuilder
        from bosdyn.client.frame_helpers import (
            ODOM_FRAME_NAME, BODY_FRAME_NAME, get_se2_a_tform_b)
        from bosdyn.client import math_helpers

        state = self.state_client.get_robot_state()
        transforms = state.kinematic_state.transforms_snapshot
        body_in_odom = get_se2_a_tform_b(
            transforms, ODOM_FRAME_NAME, BODY_FRAME_NAME)
        goal = body_in_odom * math_helpers.SE2Pose(x=distance_m, y=0, angle=0)

        cmd = RobotCommandBuilder.synchro_se2_trajectory_point_command(
            goal_x=goal.x, goal_y=goal.y, goal_heading=goal.angle,
            frame_name=ODOM_FRAME_NAME)
        self.command_client.robot_command(cmd, end_time_secs=time.time() + 30)

    def close(self):
        self.lease_keepalive.shutdown()


# ---------- Backend selection ----------
if os.environ.get("SPOT_BACKEND") == "real":
    robot = RealSpot(
        os.environ["SPOT_IP"],
        os.environ["SPOT_USER"],
        os.environ["SPOT_PASSWORD"],
    )
    print("Backend: RealSpot")
else:
    robot = FakeSpot()
    print("Backend: FakeSpot (set SPOT_BACKEND=real to command the robot)")

ACTIONS = {
    "STOP": robot.stop,
    "SIT": robot.sit,
    "STAND": robot.stand,
    "WALK_10M": robot.walk_10m,
}


# ---------- Bridge loop ----------
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(LISTEN_ADDR)
sock.settimeout(0.5)

last_cmd_time = 0.0
last_heartbeat = time.time()
moving = False

print(f"Bridge listening on {LISTEN_ADDR}")

try:
    while True:
        # dead-man check: recogniser silent while robot moving -> STOP
        if time.time() - last_heartbeat > HEARTBEAT_TIMEOUT and moving:
            print("!! heartbeat lost -> auto STOP")
            robot.stop()
            moving = False

        try:
            data, _ = sock.recvfrom(1024)
        except socket.timeout:
            continue

        try:
            msg = json.loads(data.decode())
        except json.JSONDecodeError:
            continue                      # garbage in -> ignore

        if msg.get("cmd") == "HEARTBEAT":
            last_heartbeat = time.time()
            continue

        cmd = msg.get("cmd")
        now = time.time()

        if cmd not in WHITELIST:
            print(f"?? rejected unknown command: {cmd!r}")
            continue

        # STOP always executes immediately; others are rate-limited
        if cmd != "STOP" and now - last_cmd_time < MIN_CMD_GAP:
            print(f"-- rate-limited: {cmd}")
            continue

        print(f"[{time.strftime('%H:%M:%S')}] accepted: {cmd} "
              f"(conf {msg.get('conf')})")
        ACTIONS[cmd]()
        last_cmd_time = now
        moving = cmd == "WALK_10M"

except KeyboardInterrupt:
    print("\nBridge stopping - sending STOP")
    try:
        robot.stop()
    except Exception:
        pass
    robot.close()