import os
import socket
import sys
import threading
import time

import rclpy
from rclpy.node import Node

try:
    from pynput import keyboard
except ImportError:
    keyboard = None

if sys.platform == "win32":
    import ctypes

    user32 = ctypes.WinDLL("user32", use_last_error=True)
else:
    user32 = None

if sys.platform != "win32":
    import atexit
    import select
    import termios
    import tty
else:
    atexit = None
    select = None
    termios = None
    tty = None


KEY_UP = "up"
KEY_DOWN = "down"
KEY_LEFT = "left"
KEY_RIGHT = "right"

WIN_KEY_CODES = {
    "w": 0x57,
    "a": 0x41,
    "s": 0x53,
    "d": 0x44,
    "q": 0x51,
    "e": 0x45,
    "2": 0x32,
    "3": 0x33,
    "4": 0x34,
    "5": 0x35,
    "6": 0x36,
    "up": 0x26,
    "down": 0x28,
    "left": 0x25,
    "right": 0x27,
}


class KeyboardUdpControlNode(Node):
    """ROS 2 node that sends chassis and arm UDP commands while keys are held."""

    def __init__(self):
        super().__init__("keyboard_udp_control")

        self.declare_parameter("udp_ip", "192.168.4.1")
        self.declare_parameter("udp_port", 8888)
        self.declare_parameter("chassis_linear_speed", 100)
        self.declare_parameter("chassis_angular_speed", 45)
        self.declare_parameter("arm_height_step", 5)
        self.declare_parameter("arm_joint_step", 5)
        self.declare_parameter("update_rate", 0.05)
        self.declare_parameter("debug_keys", False)
        self.declare_parameter("keyboard_backend", "auto")
        self.declare_parameter("stdin_hold_time", 0.04)

        self.udp_ip = self.get_parameter("udp_ip").value
        self.udp_port = int(self.get_parameter("udp_port").value)
        self.chassis_linear_speed = int(
            self.get_parameter("chassis_linear_speed").value
        )
        self.chassis_angular_speed = int(
            self.get_parameter("chassis_angular_speed").value
        )
        self.arm_height_step = int(self.get_parameter("arm_height_step").value)
        self.arm_joint_step = int(self.get_parameter("arm_joint_step").value)
        self.update_rate = float(self.get_parameter("update_rate").value)
        self.debug_keys = bool(self.get_parameter("debug_keys").value)
        self.keyboard_backend = str(self.get_parameter("keyboard_backend").value)
        self.stdin_hold_time = float(self.get_parameter("stdin_hold_time").value)
        self.use_windows_polling = self.should_use_windows_polling()
        self.use_stdin_keyboard = self.should_use_stdin_keyboard()

        self.chassis_x = 0
        self.chassis_y = 0
        self.chassis_w = 0

        self.arm_height = -10
        self.arm_joints = [0] * 5
        self.arm_selected_joint = 0
        self.arm_active = False

        self.pressed_keys = set()
        self.stdin_key_times = {}
        self.key_lock = threading.Lock()
        self.old_terminal_settings = None

        self._stdin_buf = ""
        self._tick_count = 0

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

        self.listener = None
        if self.use_windows_polling:
            self.get_logger().info("Using Windows keyboard polling backend.")
        elif self.use_stdin_keyboard:
            self.enable_stdin_keyboard()
            self.get_logger().info("Using stdin keyboard backend.")
        else:
            if keyboard is None:
                raise RuntimeError("pynput is required for keyboard_backend:=pynput.")
            self.listener = keyboard.Listener(
                on_press=self.on_press,
                on_release=self.on_release,
            )
            self.listener.start()
            self.get_logger().info("Using pynput keyboard listener backend.")

        self.timer = self.create_timer(self.update_rate, self.control_tick)

        self.print_help()
        self.get_logger().info(
            f"UDP target: {self.udp_ip}:{self.udp_port}"
        )
        self.get_logger().info(
            "Keyboard UDP control node started. Press Ctrl+C to exit."
        )

    def on_press(self, key):
        try:
            value = key.char.lower() if key.char else key.char
        except AttributeError:
            value = key

        with self.key_lock:
            self.pressed_keys.add(value)

    def on_release(self, key):
        try:
            value = key.char.lower() if key.char else key.char
        except AttributeError:
            value = key

        with self.key_lock:
            self.pressed_keys.discard(value)

    def should_use_windows_polling(self):
        if self.keyboard_backend == "win_poll":
            return True
        if self.keyboard_backend in ("pynput", "stdin"):
            return False
        return sys.platform == "win32" and user32 is not None

    def should_use_stdin_keyboard(self):
        if self.keyboard_backend == "stdin":
            return True
        if self.keyboard_backend in ("pynput", "win_poll"):
            return False
        return sys.platform != "win32"

    def enable_stdin_keyboard(self):
        if not sys.stdin.isatty():
            raise RuntimeError("stdin keyboard backend requires an interactive terminal.")

        self.old_terminal_settings = termios.tcgetattr(sys.stdin)
        tty.setcbreak(sys.stdin.fileno())
        atexit.register(self.restore_terminal)

    def get_pressed_keys(self):
        if self.use_windows_polling:
            return self.get_windows_pressed_keys()
        if self.use_stdin_keyboard:
            return self.get_stdin_pressed_keys()

        with self.key_lock:
            return set(self.pressed_keys)

    def get_windows_pressed_keys(self):
        keys = set()
        for key_name, key_code in WIN_KEY_CODES.items():
            if user32.GetAsyncKeyState(key_code) & 0x8000:
                keys.add(key_name)

        return keys

    _ALLOWED_SINGLE_KEYS = frozenset(
        {"w", "a", "s", "d", "q", "e", "2", "3", "4", "5", "6"}
    )
    _ARROW_MAP = {"A": KEY_UP, "B": KEY_DOWN, "C": KEY_RIGHT, "D": KEY_LEFT}

    def get_stdin_pressed_keys(self):
        now = time.monotonic()
        fd = sys.stdin.fileno()

        raw = b""
        while select.select([sys.stdin], [], [], 0)[0]:
            chunk = os.read(fd, 4096)
            if not chunk:
                break
            raw += chunk

        buf = self._stdin_buf + raw.decode("utf-8", errors="replace")

        i = 0
        n = len(buf)
        while i < n:
            ch = buf[i]

            if ch == "\x03":
                raise KeyboardInterrupt

            if ch == "\x1b":
                # Only handle when we have the full 3-byte escape.
                if i + 2 < n and buf[i + 1] in ("[", "O"):
                    arrow = buf[i + 2]
                    key = self._ARROW_MAP.get(arrow)
                    if key is not None:
                        self.stdin_key_times[key] = now
                        i += 3
                        continue
                # Incomplete / unknown escape — leave remainder for next tick.
                break

            ch_lower = ch.lower()
            if ch_lower in self._ALLOWED_SINGLE_KEYS:
                self.stdin_key_times[ch_lower] = now
            i += 1

        self._stdin_buf = buf[i:]

        expired_keys = [
            key
            for key, last_time in self.stdin_key_times.items()
            if now - last_time > self.stdin_hold_time
        ]
        for key in expired_keys:
            del self.stdin_key_times[key]

        return set(self.stdin_key_times.keys())

    def is_arm_key_pressed(self, keys):
        return (
            KEY_UP in keys
            or KEY_DOWN in keys
            or KEY_LEFT in keys
            or KEY_RIGHT in keys
            or any(str(i) in keys for i in range(2, 7))
        )

    def update_state(self, keys):
        self.chassis_x = 0
        self.chassis_y = 0
        self.chassis_w = 0

        if "w" in keys:
            self.chassis_y = -self.chassis_linear_speed
        if "s" in keys:
            self.chassis_y = self.chassis_linear_speed
        if "d" in keys:
            self.chassis_x = self.chassis_linear_speed
        if "a" in keys:
            self.chassis_x = -self.chassis_linear_speed
        if "e" in keys:
            self.chassis_w = self.chassis_angular_speed
        if "q" in keys:
            self.chassis_w = -self.chassis_angular_speed

        arm_changed = False

        for i in range(2, 7):
            if str(i) in keys:
                self.arm_selected_joint = i - 2
                arm_changed = True

        if KEY_UP in keys:
            self.arm_height = min(self.arm_height + self.arm_height_step, -10)
            arm_changed = True
        if KEY_DOWN in keys:
            self.arm_height = max(self.arm_height - self.arm_height_step, -285)
            arm_changed = True

        joint_index = self.arm_selected_joint
        if KEY_RIGHT in keys:
            self.arm_joints[joint_index] = min(
                self.arm_joints[joint_index] + self.arm_joint_step,
                180,
            )
            arm_changed = True
        if KEY_LEFT in keys:
            self.arm_joints[joint_index] = max(
                self.arm_joints[joint_index] - self.arm_joint_step,
                -180,
            )
            arm_changed = True

        if arm_changed:
            self.arm_active = True

    def build_chassis_cmd(self):
        return f"XYW:{self.chassis_x}:{self.chassis_y}:{self.chassis_w}:XZHY\n".encode()

    def build_arm_cmd(self):
        joints = self.arm_joints
        return (
            f"JXB:{self.arm_height}:"
            f"{joints[0]}:{joints[1]}:{joints[2]}:"
            f"{joints[3]}:{joints[4]}:0:0:EZHY\n"
        ).encode()

    def control_tick(self):
        keys = self.get_pressed_keys()
        self.update_state(keys)

        self.sock.sendto(
            self.build_chassis_cmd(),
            (self.udp_ip, self.udp_port),
        )

        if self.arm_active or self.is_arm_key_pressed(keys):
            self.sock.sendto(self.build_arm_cmd(), (self.udp_ip, self.udp_port))

        self._tick_count += 1
        if self._tick_count % 5 == 0:
            self.print_status(keys)

    def print_status(self, keys):
        selected_joint = self.arm_selected_joint + 2
        joints = self.arm_joints
        key_text = ""
        if self.debug_keys:
            key_text = "  Keys:" + ",".join(self.format_key(key) for key in keys)

        status = (
            f"\rChassis X:{self.chassis_x:>4}  Y:{self.chassis_y:>4}  "
            f"W:{self.chassis_w:>4}  |  "
            f"Arm H:{self.arm_height:>4}mm  J{selected_joint}:"
            f"[{joints[0]:>4} {joints[1]:>4} {joints[2]:>4} "
            f"{joints[3]:>4} {joints[4]:>4}]  "
            f"{'[ARM]' if self.arm_active else '[STANDBY]'}"
            f"{key_text}"
        )
        print(status, end="", flush=True)

    def format_key(self, key):
        if isinstance(key, str):
            return key
        return str(key).replace("Key.", "")

    def send_stop(self):
        self.sock.sendto(
            b"XYW:0:0:0:XZHY\n",
            (self.udp_ip, self.udp_port),
        )

    def destroy_node(self):
        self.send_stop()

        if self.listener is not None and self.listener.running:
            self.listener.stop()
        self.restore_terminal()

        self.sock.close()

        print("\nSent chassis stop command.")
        super().destroy_node()

    def restore_terminal(self):
        if self.old_terminal_settings is not None:
            termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_terminal_settings)
            self.old_terminal_settings = None

    def print_help(self):
        print("=" * 62)
        print("  ROS 2 chassis + arm keyboard UDP control")
        print("=" * 62)
        print("  Chassis:")
        print("    W/S    forward / backward")
        print("    A/D    left / right")
        print("    Q/E    turn left / turn right")
        print("  Arm:")
        print("    Up/Down  lift up / down")
        print("    2 ~ 6    select joint J2 ~ J6")
        print("    Lt/Rt    decrease / increase selected joint angle")
        print("  Other:")
        print("    Ctrl+C   exit and send chassis stop command")
        print("=" * 62)
        print("  Arm commands are sent only after pressing an arm key.")
        print("  Make sure the arm pose matches the initial values.")
        print("=" * 62)


def main(args=None):
    rclpy.init(args=args)
    node = KeyboardUdpControlNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
