"""
gesture_recognition_node.py — Hand and body gesture recognition via MediaPipe.

Uses MediaPipe Hands and Pose to detect gestures on Jetson Orin GPU.

Recognizes:
  Hand gestures: wave, point, stop_palm (e-stop), thumbs_up, come_here (beckon)
  Body gestures: arms_up (stop), arms_spread (back off)

Publishes:
  /saltybot/gestures   saltybot_social_msgs/GestureArray   10+ fps

Parameters:
  camera_topic         str     '/camera/color/image_raw'    RGB camera input
  confidence_threshold float   0.7                           detection confidence
  publish_hz           float   15.0                          output rate (10+ fps target)
  max_distance_m       float   5.0                           max gesture range
  enable_gpu           bool    true                          use GPU acceleration
"""

from __future__ import annotations

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

import numpy as np
import cv2
from cv_bridge import CvBridge
import threading
from collections import deque
from typing import Optional

from std_msgs.msg import Header
from sensor_msgs.msg import Image
from geometry_msgs.msg import Point

try:
    from saltybot_social_msgs.msg import Gesture, GestureArray
    _GESTURE_MSGS_OK = True
except ImportError:
    _GESTURE_MSGS_OK = False

try:
    import mediapipe as mp
    _MEDIAPIPE_OK = True
except ImportError:
    _MEDIAPIPE_OK = False


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


class GestureDetector:
    """MediaPipe-based gesture detector for hands and pose."""

    # Hand gesture thresholds
    GESTURE_DISTANCE_THRESHOLD = 0.05
    WAVE_DURATION = 5  # frames
    BECKON_DURATION = 4
    POINT_MIN_EXTEND = 0.3  # index extension threshold

    def __init__(self, enable_gpu: bool = True):
        if not _MEDIAPIPE_OK:
            raise ImportError("MediaPipe not available")

        self.enable_gpu = enable_gpu

        # Initialize MediaPipe
        self.mp_hands = mp.solutions.hands
        self.mp_pose = mp.solutions.pose
        self.mp_drawing = mp.solutions.drawing_utils

        # Create hand detector
        self.hands = self.mp_hands.Hands(
            static_image_mode=False,
            max_num_hands=10,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5,
            model_complexity=0,  # 0=lite (faster), 1=full
        )

        # Create pose detector
        self.pose = self.mp_pose.Pose(
            static_image_mode=False,
            model_complexity=1,
            smooth_landmarks=True,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5,
        )

        # Gesture history for temporal smoothing
        self.hand_history = deque(maxlen=10)
        self.pose_history = deque(maxlen=10)

    def detect_hand_gestures(self, frame: np.ndarray, person_id: int = -1) -> list[dict]:
        """
        Detect hand gestures using MediaPipe Hands.

        Returns:
            List of detected gestures with type, confidence, position
        """
        gestures = []

        if frame is None or frame.size == 0:
            return gestures

        try:
            # Convert BGR to RGB
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, _ = rgb_frame.shape

            # Detect hands
            results = self.hands.process(rgb_frame)

            if not results.multi_hand_landmarks or not results.multi_handedness:
                return gestures

            for hand_landmarks, handedness in zip(
                results.multi_hand_landmarks, results.multi_handedness
            ):
                is_right = handedness.classification[0].label == "Right"
                confidence = handedness.classification[0].score

                # Extract key landmarks
                landmarks = np.array(
                    [[lm.x, lm.y, lm.z] for lm in hand_landmarks.landmark]
                )

                # Detect specific hand gestures
                gesture_type, gesture_conf = self._classify_hand_gesture(
                    landmarks, is_right
                )

                if gesture_type:
                    # Get hand center position
                    hand_x = float(np.mean(landmarks[:, 0]))
                    hand_y = float(np.mean(landmarks[:, 1]))

                    gestures.append({
                        'type': gesture_type,
                        'confidence': float(gesture_conf * confidence),
                        'hand_x': hand_x,
                        'hand_y': hand_y,
                        'is_right_hand': is_right,
                        'source': 'hand',
                        'person_id': person_id,
                    })

            self.hand_history.append(gestures)

        except Exception as e:
            pass

        return gestures

    def detect_body_gestures(self, frame: np.ndarray, person_id: int = -1) -> list[dict]:
        """
        Detect body/pose gestures using MediaPipe Pose.

        Returns:
            List of detected pose-based gestures
        """
        gestures = []

        if frame is None or frame.size == 0:
            return gestures

        try:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, _ = rgb_frame.shape

            results = self.pose.process(rgb_frame)

            if not results.pose_landmarks:
                return gestures

            landmarks = np.array(
                [[lm.x, lm.y, lm.z] for lm in results.pose_landmarks.landmark]
            )

            # Detect specific body gestures
            gesture_type, gesture_conf = self._classify_body_gesture(landmarks)

            if gesture_type:
                # Get body center
                body_x = float(np.mean(landmarks[:, 0]))
                body_y = float(np.mean(landmarks[:, 1]))

                gestures.append({
                    'type': gesture_type,
                    'confidence': float(gesture_conf),
                    'hand_x': body_x,
                    'hand_y': body_y,
                    'is_right_hand': False,
                    'source': 'body_pose',
                    'person_id': person_id,
                })

            self.pose_history.append(gestures)

        except Exception as e:
            pass

        return gestures

    def _classify_hand_gesture(
        self, landmarks: np.ndarray, is_right: bool
    ) -> tuple[Optional[str], float]:
        """
        Classify hand gesture from MediaPipe landmarks.

        Returns:
            (gesture_type, confidence)
        """
        if landmarks.shape[0] < 21:
            return None, 0.0

        # Landmark indices
        # 0: wrist, 5: index, 9: middle, 13: ring, 17: pinky
        # 4: thumb tip, 8: index tip, 12: middle tip, 16: ring tip, 20: pinky tip

        wrist = landmarks[0]
        thumb_tip = landmarks[4]
        index_tip = landmarks[8]
        middle_tip = landmarks[12]
        ring_tip = landmarks[16]
        pinky_tip = landmarks[20]

        # Palm normal (pointing direction)
        palm_normal = self._get_palm_normal(landmarks)

        # Finger extension
        index_extended = self._distance(index_tip, landmarks[5]) > self.POINT_MIN_EXTEND
        middle_extended = self._distance(middle_tip, landmarks[9]) > self.POINT_MIN_EXTEND
        ring_extended = self._distance(ring_tip, landmarks[13]) > self.POINT_MIN_EXTEND
        pinky_extended = self._distance(pinky_tip, landmarks[17]) > self.POINT_MIN_EXTEND
        thumb_extended = self._distance(thumb_tip, landmarks[2]) > 0.1

        # Thumbs-up: thumb extended up, hand vertical
        if thumb_extended and not (index_extended or middle_extended):
            palm_y = np.mean([landmarks[i][1] for i in [5, 9, 13, 17]])
            if thumb_tip[1] < palm_y - 0.1:  # Thumb above palm
                return 'thumbs_up', 0.85

        # Stop palm: all fingers extended, palm forward
        if index_extended and middle_extended and ring_extended and pinky_extended:
            if palm_normal[2] > 0.3:  # Palm facing camera
                return 'stop_palm', 0.8

        # Point: only index extended
        if index_extended and not (middle_extended or ring_extended or pinky_extended):
            return 'point', 0.8

        # Wave: hand moving (approximate via history)
        if len(self.hand_history) > self.WAVE_DURATION:
            if self._detect_wave_motion():
                return 'wave', 0.75

        # Come-here (beckon): curled fingers, repetitive motion
        if not (index_extended or middle_extended):
            if len(self.hand_history) > self.BECKON_DURATION:
                if self._detect_beckon_motion():
                    return 'come_here', 0.75

        return None, 0.0

    def _classify_body_gesture(self, landmarks: np.ndarray) -> tuple[Optional[str], float]:
        """
        Classify body gesture from MediaPipe Pose landmarks.

        Returns:
            (gesture_type, confidence)
        """
        if landmarks.shape[0] < 33:
            return None, 0.0

        # Key body landmarks
        left_shoulder = landmarks[11]
        right_shoulder = landmarks[12]
        left_hip = landmarks[23]
        right_hip = landmarks[24]
        left_wrist = landmarks[9]
        right_wrist = landmarks[10]

        shoulder_y = np.mean([left_shoulder[1], right_shoulder[1]])
        hip_y = np.mean([left_hip[1], right_hip[1]])
        wrist_y_max = max(left_wrist[1], right_wrist[1])

        # Arms up (emergency stop)
        if wrist_y_max < shoulder_y - 0.2:
            return 'arms_up', 0.85

        # Arms spread (back off)
        shoulder_dist = self._distance(left_shoulder[:2], right_shoulder[:2])
        wrist_dist = self._distance(left_wrist[:2], right_wrist[:2])
        if wrist_dist > shoulder_dist * 1.2:
            return 'arms_spread', 0.8

        # Crouch (come closer)
        if hip_y - shoulder_y > 0.3:
            return 'crouch', 0.8

        return None, 0.0

    def _get_palm_normal(self, landmarks: np.ndarray) -> np.ndarray:
        """Compute palm normal vector (pointing direction)."""
        wrist = landmarks[0]
        middle_mcp = landmarks[9]
        index_mcp = landmarks[5]
        v1 = index_mcp - wrist
        v2 = middle_mcp - wrist
        normal = np.cross(v1, v2)
        return normal / (np.linalg.norm(normal) + 1e-6)

    def _distance(self, p1: np.ndarray, p2: np.ndarray) -> float:
        """Euclidean distance between two points."""
        return float(np.linalg.norm(p1 - p2))

    def _detect_wave_motion(self) -> bool:
        """Detect waving motion from hand history."""
        if len(self.hand_history) < self.WAVE_DURATION:
            return False
        # Simple heuristic: high variance in x-position over time
        x_positions = [g[0]['hand_x'] for g in self.hand_history if g]
        if len(x_positions) < self.WAVE_DURATION:
            return False
        return float(np.std(x_positions)) > 0.05

    def _detect_beckon_motion(self) -> bool:
        """Detect beckoning motion from hand history."""
        if len(self.hand_history) < self.BECKON_DURATION:
            return False
        # High variance in y-position (up-down motion)
        y_positions = [g[0]['hand_y'] for g in self.hand_history if g]
        if len(y_positions) < self.BECKON_DURATION:
            return False
        return float(np.std(y_positions)) > 0.04


class GestureRecognitionNode(Node):

    def __init__(self):
        super().__init__('gesture_recognition')

        # Parameters
        self.declare_parameter('camera_topic', '/camera/color/image_raw')
        self.declare_parameter('confidence_threshold', 0.7)
        self.declare_parameter('publish_hz', 15.0)
        self.declare_parameter('max_distance_m', 5.0)
        self.declare_parameter('enable_gpu', True)

        camera_topic = self.get_parameter('camera_topic').value
        self.confidence_threshold = self.get_parameter('confidence_threshold').value
        pub_hz = self.get_parameter('publish_hz').value
        max_distance = self.get_parameter('max_distance_m').value
        enable_gpu = self.get_parameter('enable_gpu').value

        # Publisher
        self._pub_gestures = None
        if _GESTURE_MSGS_OK:
            self._pub_gestures = self.create_publisher(
                GestureArray, '/saltybot/gestures', 10, qos_profile=_SENSOR_QOS
            )
        else:
            self.get_logger().error('saltybot_social_msgs not available')
            return

        # Gesture detector
        self._detector: Optional[GestureDetector] = None
        self._detector_lock = threading.Lock()

        if _MEDIAPIPE_OK:
            try:
                self._detector = GestureDetector(enable_gpu=enable_gpu)
            except Exception as e:
                self.get_logger().error(f'Failed to initialize MediaPipe: {e}')

        # Video bridge
        self._bridge = CvBridge()
        self._latest_image: Image | None = None
        self._lock = threading.Lock()

        # Subscriptions
        self.create_subscription(Image, camera_topic, self._on_image, _SENSOR_QOS)

        # Publish timer
        self.create_timer(1.0 / pub_hz, self._tick)

        self.get_logger().info(
            f'gesture_recognition ready — '
            f'camera={camera_topic} confidence_threshold={self.confidence_threshold} hz={pub_hz}'
        )

    def _on_image(self, msg: Image) -> None:
        with self._lock:
            self._latest_image = msg

    def _tick(self) -> None:
        """Detect and publish gestures."""
        if self._pub_gestures is None or self._detector is None:
            return

        with self._lock:
            if self._latest_image is None:
                return
            image_msg = self._latest_image

        try:
            frame = self._bridge.imgmsg_to_cv2(
                image_msg, desired_encoding='bgr8'
            )
        except Exception as e:
            self.get_logger().warn(f'Image conversion error: {e}')
            return

        # Detect hand and body gestures
        hand_gestures = self._detector.detect_hand_gestures(frame)
        body_gestures = self._detector.detect_body_gestures(frame)

        all_gestures = hand_gestures + body_gestures

        # Filter by confidence threshold
        filtered_gestures = [
            g for g in all_gestures if g['confidence'] >= self.confidence_threshold
        ]

        # Build and publish GestureArray
        gesture_array = GestureArray()
        gesture_array.header = Header(
            stamp=self.get_clock().now().to_msg(),
            frame_id='camera',
        )

        for g in filtered_gestures:
            gesture = Gesture()
            gesture.header = gesture_array.header
            gesture.gesture_type = g['type']
            gesture.person_id = g.get('person_id', -1)
            gesture.confidence = g['confidence']
            gesture.hand_x = g['hand_x']
            gesture.hand_y = g['hand_y']
            gesture.is_right_hand = g['is_right_hand']
            gesture.source = g['source']

            # Map point direction if applicable
            if g['type'] == 'point':
                if g['hand_x'] < 0.33:
                    gesture.direction = 'left'
                elif g['hand_x'] > 0.67:
                    gesture.direction = 'right'
                elif g['hand_y'] < 0.33:
                    gesture.direction = 'up'
                else:
                    gesture.direction = 'forward'

            gesture_array.gestures.append(gesture)

        gesture_array.count = len(gesture_array.gestures)
        self._pub_gestures.publish(gesture_array)


def main(args=None):
    rclpy.init(args=args)
    node = GestureRecognitionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
