"""
hand_tracking_node.py — MediaPipe Hands inference node (Issue #342).

Subscribes
----------
  /camera/color/image_raw  (sensor_msgs/Image)

Publishes
---------
  /saltybot/hands          (saltybot_hand_tracking_msgs/HandLandmarksArray)
  /saltybot/hand_gesture   (std_msgs/String)

Parameters
----------
  max_hands          int   2      Maximum hands detected per frame
  model_complexity   int   0      MediaPipe model: 0=lite, 1=full (0 for 20+ FPS)
  min_detection_conf float 0.60   MediaPipe detection confidence threshold
  min_tracking_conf  float 0.50   MediaPipe tracking confidence threshold
  gesture_min_conf   float 0.60   Minimum gesture confidence to publish on hand_gesture
  wave_history_len   int   24     Frames kept in WaveDetector history
  wave_min_reversals int   2      Oscillation reversals needed for wave
  wave_min_amplitude float 0.08   Peak-to-peak wrist-x amplitude for wave

Performance note
----------------
  model_complexity=0 (lite) on Orin Nano Super (1024-core Ampere, 67 TOPS)
  targets >20 FPS at 640×480.  Drop to 480×360 in the camera launch if needed.
"""

from __future__ import annotations

import threading
import time
from typing import Dict, List, Optional, Tuple

import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSHistoryPolicy

from sensor_msgs.msg import Image
from std_msgs.msg import String

from saltybot_hand_tracking_msgs.msg import HandLandmarks, HandLandmarksArray

from ._hand_gestures import Landmark, WaveDetector, classify_hand

# Optional runtime imports — guarded so unit tests don't need them
try:
    import cv2
    _HAS_CV = True
except ImportError:
    _HAS_CV = False

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


# ── MediaPipe wrapper ──────────────────────────────────────────────────────────

class _MPHands:
    """Thin wrapper around mediapipe.solutions.hands with lazy init."""

    def __init__(self, max_hands: int, complexity: int,
                 det_conf: float, trk_conf: float) -> None:
        self._max_hands  = max_hands
        self._complexity = complexity
        self._det_conf   = det_conf
        self._trk_conf   = trk_conf
        self._hands      = None

    def init(self) -> None:
        if not _HAS_MP:
            return
        self._hands = mp.solutions.hands.Hands(
            static_image_mode=False,
            max_num_hands=self._max_hands,
            min_detection_confidence=self._det_conf,
            min_tracking_confidence=self._trk_conf,
            model_complexity=self._complexity,
        )

    def process(self, bgr: np.ndarray):
        """Process a BGR image; returns mediapipe Hands results or None."""
        if self._hands is None or not _HAS_MP or not _HAS_CV:
            return None
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        rgb.flags.writeable = False
        try:
            return self._hands.process(rgb)
        except Exception:
            return None

    def close(self) -> None:
        if self._hands is not None:
            self._hands.close()


# ── ROS2 Node ─────────────────────────────────────────────────────────────────

class HandTrackingNode(Node):
    def __init__(self) -> None:
        super().__init__('hand_tracking_node')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('max_hands',          2)
        self.declare_parameter('model_complexity',   0)
        self.declare_parameter('min_detection_conf', 0.60)
        self.declare_parameter('min_tracking_conf',  0.50)
        self.declare_parameter('gesture_min_conf',   0.60)
        self.declare_parameter('wave_history_len',   24)
        self.declare_parameter('wave_min_reversals', 2)
        self.declare_parameter('wave_min_amplitude', 0.08)

        p = self.get_parameter
        self._gesture_min_conf: float = p('gesture_min_conf').value

        # ── QoS ─────────────────────────────────────────────────────────────
        img_qos = QoSProfile(
            reliability=QoSReliabilityPolicy.BEST_EFFORT,
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # ── Publishers ───────────────────────────────────────────────────────
        self._hands_pub = self.create_publisher(
            HandLandmarksArray, '/saltybot/hands', 10
        )
        self._gesture_pub = self.create_publisher(
            String, '/saltybot/hand_gesture', 10
        )

        # ── Subscriber ───────────────────────────────────────────────────────
        self._sub = self.create_subscription(
            Image, '/camera/color/image_raw',
            self._image_cb, img_qos,
        )

        # ── Per-hand WaveDetectors keyed by hand index ────────────────────
        self._wave_dets: Dict[int, WaveDetector] = {}
        wave_hist = p('wave_history_len').value
        wave_rev  = p('wave_min_reversals').value
        wave_amp  = p('wave_min_amplitude').value
        self._wave_kwargs = dict(
            history_len=wave_hist,
            min_reversals=wave_rev,
            min_amplitude=wave_amp,
        )

        # ── MediaPipe (background init) ───────────────────────────────────
        self._mp = _MPHands(
            max_hands  = p('max_hands').value,
            complexity = p('model_complexity').value,
            det_conf   = p('min_detection_conf').value,
            trk_conf   = p('min_tracking_conf').value,
        )
        self._mp_ready = threading.Event()
        threading.Thread(target=self._init_mp, daemon=True).start()

        self.get_logger().info(
            f'hand_tracking_node ready — '
            f'max_hands={p("max_hands").value}, '
            f'complexity={p("model_complexity").value}'
        )

    # ── Init ──────────────────────────────────────────────────────────────────

    def _init_mp(self) -> None:
        if not _HAS_MP:
            self.get_logger().warn(
                'mediapipe not installed — no hand tracking. '
                'Install: pip install mediapipe'
            )
            return
        t0 = time.time()
        self._mp.init()
        self._mp_ready.set()
        self.get_logger().info(
            f'MediaPipe Hands ready ({time.time() - t0:.1f}s)'
        )

    # ── Image callback ────────────────────────────────────────────────────────

    def _image_cb(self, msg: Image) -> None:
        if not self._mp_ready.is_set():
            return

        bgr = self._ros_to_bgr(msg)
        if bgr is None:
            return

        results = self._mp.process(bgr)

        arr            = HandLandmarksArray()
        arr.header     = msg.header
        arr.hand_count = 0
        best_gesture   = ""
        best_conf      = 0.0

        if results and results.multi_hand_landmarks:
            for hand_idx, (hand_lm, hand_info) in enumerate(
                zip(results.multi_hand_landmarks, results.multi_handedness)
            ):
                cls       = hand_info.classification[0]
                is_right  = cls.label == "Right"
                hs_score  = float(cls.score)

                lm = [Landmark(p.x, p.y, p.z) for p in hand_lm.landmark]

                wave_det = self._wave_dets.get(hand_idx)
                if wave_det is None:
                    wave_det = WaveDetector(**self._wave_kwargs)
                    self._wave_dets[hand_idx] = wave_det

                gesture_result = classify_hand(lm, is_right=is_right, wave_det=wave_det)

                # Build HandLandmarks message
                hl              = HandLandmarks()
                hl.header       = msg.header
                hl.is_right_hand     = is_right
                hl.handedness_score  = hs_score
                hl.landmark_xyz      = self._pack_landmarks(lm)
                hl.gesture           = gesture_result.gesture
                hl.point_direction   = gesture_result.direction
                hl.gesture_confidence = float(gesture_result.confidence)
                hl.wrist_x           = float(lm[0].x)
                hl.wrist_y           = float(lm[0].y)

                arr.hands.append(hl)
                arr.hand_count += 1

                if gesture_result.confidence > best_conf:
                    best_conf    = gesture_result.confidence
                    best_gesture = gesture_result.gesture

        self._hands_pub.publish(arr)

        # Publish hand_gesture String only when confident enough
        if best_gesture and best_gesture != "none" \
                and best_conf >= self._gesture_min_conf:
            gs = String()
            gs.data = best_gesture
            self._gesture_pub.publish(gs)

    # ── Helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _pack_landmarks(lm: List[Landmark]) -> List[float]:
        """Pack 21 Landmark objects into a flat [x0,y0,z0, ..., x20,y20,z20] list."""
        out: List[float] = []
        for l in lm:
            out.extend([l.x, l.y, l.z])
        return out

    @staticmethod
    def _ros_to_bgr(msg: Image) -> Optional[np.ndarray]:
        """Convert sensor_msgs/Image to uint8 BGR numpy array."""
        enc  = msg.encoding.lower()
        data = np.frombuffer(msg.data, dtype=np.uint8)
        if enc == 'bgr8':
            try:
                return data.reshape((msg.height, msg.width, 3))
            except ValueError:
                return None
        if enc == 'rgb8':
            try:
                img = data.reshape((msg.height, msg.width, 3))
                return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) if _HAS_CV else None
            except ValueError:
                return None
        if enc == 'mono8':
            try:
                img = data.reshape((msg.height, msg.width))
                return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) if _HAS_CV else None
            except ValueError:
                return None
        return None

    # ── Cleanup ───────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        self._mp.close()
        super().destroy_node()


# ── Entry point ───────────────────────────────────────────────────────────────

def main(args=None) -> None:
    rclpy.init(args=args)
    node = HandTrackingNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == '__main__':
    main()
