"""
person_tracking_node.py — P0 follow-me person tracking node (Issue #363).

Runs a real-time person detection + tracking pipeline on the D435i colour and
depth streams and publishes a single TargetTrack message for the follow-me
motion controller.

Detection backend
-----------------
  Attempts YOLOv8n via ultralytics (auto-converted to TensorRT FP16 on first
  run for ≥ 15 fps on Orin Nano).  Falls back to a simple HOG+SVM detector
  when ultralytics is unavailable.

Subscribes
----------
  /camera/color/image_raw           sensor_msgs/Image  (BEST_EFFORT)
  /camera/depth/image_rect_raw      sensor_msgs/Image  (BEST_EFFORT)
  /camera/depth/camera_info         sensor_msgs/CameraInfo  (RELIABLE)
  /saltybot/follow_start            std_msgs/Empty  — start/resume following
  /saltybot/follow_stop             std_msgs/Empty  — stop following

Publishes
---------
  /saltybot/target_track   saltybot_scene_msgs/TargetTrack

Parameters
----------
  detector_model    str    'yolov8n.pt'   Ultralytics model file or TRT engine
  use_tensorrt      bool   True           Convert to TensorRT FP16 on first run
  max_fps           float  30.0           Maximum processing rate (Hz)
  iou_threshold     float  0.30           Tracker IoU matching threshold
  min_hits          int    3              Frames before TENTATIVE → ACTIVE
  max_lost_frames   int    30             Frames a track survives without det
  reid_threshold    float  0.55           HSV histogram re-ID similarity cutoff
  depth_scale       float  0.001          D435i raw-to-metres scale
  depth_max_m       float  5.0            Range beyond which depth degrades
  auto_follow       bool   True           Auto-select nearest person on start
"""

from __future__ import annotations

import threading
import time
from typing import List, Optional

import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy,
)

from sensor_msgs.msg import CameraInfo, Image
from std_msgs.msg import Empty
from saltybot_scene_msgs.msg import TargetTrack
from ._person_tracker import (
    BBox, CamParams, Detection,
    PersonTracker, FollowTargetSelector,
    DEPTH_GOOD, DEPTH_MARGINAL,
)


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=2,
)
_RELIABLE_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)


# ── Detector wrappers ─────────────────────────────────────────────────────────

class _YoloDetector:
    """Lazy-initialised YOLOv8n person detector (ultralytics / TensorRT)."""

    def __init__(
        self,
        model_path:   str  = 'yolov8n.pt',
        use_tensorrt: bool = True,
        logger=None,
    ) -> None:
        self._model_path   = model_path
        self._use_trt      = use_tensorrt
        self._log          = logger
        self._model        = None
        self._available    = False
        self._ready        = threading.Event()
        threading.Thread(target=self._load, daemon=True).start()

    def _load(self) -> None:
        try:
            from ultralytics import YOLO
            model = YOLO(self._model_path)
            if self._use_trt:
                try:
                    model = YOLO(model.export(format='engine', half=True, device=0))
                except Exception as e:
                    if self._log:
                        self._log.warn(f'TRT export failed ({e}); using PyTorch')
            self._model     = model
            self._available = True
            if self._log:
                self._log.info(f'YOLO detector loaded: {self._model_path}')
        except Exception as e:
            self._available = False
            if self._log:
                self._log.warn(f'ultralytics not available ({e}); using HOG fallback')
        finally:
            self._ready.set()

    def detect(self, bgr: np.ndarray, conf_thresh: float = 0.40) -> List[Detection]:
        if not self._ready.wait(timeout=0.01) or not self._available:
            return []
        results = self._model(bgr, classes=[0], conf=conf_thresh, verbose=False)
        dets: List[Detection] = []
        for r in results:
            for box in r.boxes:
                x1, y1, x2, y2 = (int(v) for v in box.xyxy[0].cpu().numpy())
                w, h = max(1, x2 - x1), max(1, y2 - y1)
                dets.append(Detection(
                    bbox       = BBox(x1, y1, w, h),
                    confidence = float(box.conf[0]),
                    frame_bgr  = bgr,
                ))
        return dets


class _HogDetector:
    """OpenCV HOG+SVM person detector — CPU fallback, ~5–10 fps."""

    def __init__(self) -> None:
        import cv2
        self._hog = cv2.HOGDescriptor()
        self._hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

    def detect(self, bgr: np.ndarray, conf_thresh: float = 0.40) -> List[Detection]:
        import cv2
        small = cv2.resize(bgr, (320, 240)) if bgr.shape[1] > 320 else bgr
        scale = bgr.shape[1] / small.shape[1]
        rects, weights = self._hog.detectMultiScale(
            small, winStride=(8, 8), padding=(4, 4), scale=1.05,
        )
        dets: List[Detection] = []
        for rect, w in zip(rects, weights):
            conf = float(np.clip(w, 0.0, 1.0))
            if conf < conf_thresh:
                continue
            x, y, rw, rh = rect
            dets.append(Detection(
                bbox       = BBox(
                    int(x * scale), int(y * scale),
                    int(rw * scale), int(rh * scale),
                ),
                confidence = conf,
                frame_bgr  = bgr,
            ))
        return dets


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

class PersonTrackingNode(Node):

    def __init__(self) -> None:
        super().__init__('person_tracking_node')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('detector_model',  'yolov8n.pt')
        self.declare_parameter('use_tensorrt',    True)
        self.declare_parameter('max_fps',         30.0)
        self.declare_parameter('iou_threshold',   0.30)
        self.declare_parameter('min_hits',        3)
        self.declare_parameter('max_lost_frames', 30)
        self.declare_parameter('reid_threshold',  0.55)
        self.declare_parameter('depth_scale',     0.001)
        self.declare_parameter('depth_max_m',     5.0)
        self.declare_parameter('auto_follow',     True)

        p = self.get_parameter
        self._min_period   = 1.0 / max(float(p('max_fps').value), 1.0)
        self._depth_scale  = float(p('depth_scale').value)
        self._depth_max_m  = float(p('depth_max_m').value)
        self._auto_follow  = bool(p('auto_follow').value)

        # ── Tracker + selector ───────────────────────────────────────────────
        self._tracker = PersonTracker(
            iou_threshold   = float(p('iou_threshold').value),
            min_hits        = int(p('min_hits').value),
            max_lost_frames = int(p('max_lost_frames').value),
            reid_threshold  = float(p('reid_threshold').value),
        )
        self._selector = FollowTargetSelector(hold_frames=15)
        if self._auto_follow:
            self._selector.start()

        # ── Detector (lazy) ──────────────────────────────────────────────────
        self._detector: Optional[_YoloDetector] = _YoloDetector(
            model_path   = str(p('detector_model').value),
            use_tensorrt = bool(p('use_tensorrt').value),
            logger       = self.get_logger(),
        )
        self._hog_fallback: Optional[_HogDetector] = None

        # ── Camera state ─────────────────────────────────────────────────────
        self._cam          = CamParams()
        self._cam_received = False
        self._depth_image: Optional[np.ndarray] = None
        self._last_proc    = 0.0

        # ── Subscribers ──────────────────────────────────────────────────────
        self._color_sub = self.create_subscription(
            Image, '/camera/color/image_raw',
            self._on_color, _SENSOR_QOS,
        )
        self._depth_sub = self.create_subscription(
            Image, '/camera/depth/image_rect_raw',
            self._on_depth, _SENSOR_QOS,
        )
        self._info_sub = self.create_subscription(
            CameraInfo, '/camera/depth/camera_info',
            self._on_cam_info, _RELIABLE_QOS,
        )
        self._start_sub = self.create_subscription(
            Empty, '/saltybot/follow_start',
            lambda _: self._selector.start(), 10,
        )
        self._stop_sub = self.create_subscription(
            Empty, '/saltybot/follow_stop',
            lambda _: self._selector.stop(), 10,
        )

        # ── Publisher ────────────────────────────────────────────────────────
        self._pub = self.create_publisher(
            TargetTrack, '/saltybot/target_track', 10
        )

        self.get_logger().info(
            'person_tracking_node ready — '
            f'auto_follow={self._auto_follow}, '
            f'max_fps={1.0/self._min_period:.0f}'
        )

    # ── Callbacks ─────────────────────────────────────────────────────────────

    def _on_cam_info(self, msg: CameraInfo) -> None:
        if self._cam_received:
            return
        self._cam = CamParams(
            fx=float(msg.k[0]), fy=float(msg.k[4]),
            cx=float(msg.k[2]), cy=float(msg.k[5]),
        )
        self._cam_received = True
        self.get_logger().info(
            f'Camera intrinsics: fx={self._cam.fx:.1f} cx={self._cam.cx:.1f}'
        )

    def _on_depth(self, msg: Image) -> None:
        if msg.encoding.lower() not in ('16uc1', 'mono16'):
            return
        data = np.frombuffer(msg.data, dtype=np.uint16)
        try:
            self._depth_image = data.reshape((msg.height, msg.width))
        except ValueError:
            pass

    def _on_color(self, msg: Image) -> None:
        now = time.monotonic()
        if now - self._last_proc < self._min_period:
            return
        self._last_proc = now

        enc = msg.encoding.lower()
        if enc not in ('bgr8', 'rgb8'):
            return
        data = np.frombuffer(msg.data, dtype=np.uint8)
        try:
            frame = data.reshape((msg.height, msg.width, 3))
        except ValueError:
            return
        if enc == 'rgb8':
            import cv2
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

        # ── Detect ────────────────────────────────────────────────────────
        dets: List[Detection] = []
        if self._detector and self._detector._ready.is_set():
            if self._detector._available:
                dets = self._detector.detect(frame)
            else:
                # Init HOG fallback on first YOLO failure
                if self._hog_fallback is None:
                    try:
                        self._hog_fallback = _HogDetector()
                    except Exception:
                        pass
                if self._hog_fallback:
                    dets = self._hog_fallback.detect(frame)
        else:
            # YOLO still loading — run HOG if available
            if self._hog_fallback is None:
                try:
                    self._hog_fallback = _HogDetector()
                except Exception:
                    pass
            if self._hog_fallback:
                dets = self._hog_fallback.detect(frame)

        # ── Track ─────────────────────────────────────────────────────────
        active = self._tracker.update(
            dets,
            cam        = self._cam if self._cam_received else None,
            depth_u16  = self._depth_image,
            depth_scale = self._depth_scale,
        )

        # ── Select target ─────────────────────────────────────────────────
        target = self._selector.update(
            active, img_cx=self._cam.cx
        )

        # ── Publish ───────────────────────────────────────────────────────
        out = TargetTrack()
        out.header.stamp    = msg.header.stamp
        out.header.frame_id = msg.header.frame_id or 'camera_color_optical_frame'

        if target is not None:
            out.tracking_active = True
            out.track_id        = target.track_id
            out.bearing_deg     = float(target.bearing)
            out.distance_m      = float(target.distance)
            out.confidence      = float(target.confidence)
            out.bbox_x          = int(target.bbox.x)
            out.bbox_y          = int(target.bbox.y)
            out.bbox_w          = int(target.bbox.w)
            out.bbox_h          = int(target.bbox.h)
            out.depth_quality   = int(target.depth_qual)

            # Convert Kalman pixel velocity → bearing rate
            if self._cam_received and target.vel_u != 0.0:
                u_c = target.bbox.x + target.bbox.w * 0.5
                # d(bearing)/du ≈ fx / (fx² + (u-cx)²) * (180/π)
                denom = self._cam.fx ** 2 + (u_c - self._cam.cx) ** 2
                d_bear_du = (self._cam.fx / denom) * (180.0 / 3.14159)
                out.vel_bearing_dps = float(d_bear_du * target.vel_u * self._cam.fps)

            # Distance velocity from depth (placeholder: not computed per-frame here)
            out.vel_dist_mps = 0.0
        else:
            out.tracking_active = False

        self._pub.publish(out)


def main(args=None) -> None:
    rclpy.init(args=args)
    node = PersonTrackingNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
