"""
person_reid_node.py — Appearance-based person re-identification (Issue #322).

Subscribes to a camera RGB stream and person detection bounding boxes,
extracts per-person HSV colour histograms, matches against an in-memory
gallery using cosine similarity, and publishes stable PersonTrack IDs.

Running one instance per camera (remapping /camera/color/image_raw and
/saltybot/scene/objects) achieves cross-camera re-id: the same physical
person is assigned the same track_id regardless of which camera view
they appear in, because appearance matching is camera-agnostic.

Subscribes
----------
  /camera/color/image_raw       sensor_msgs/Image              (BEST_EFFORT)
  /saltybot/scene/objects        saltybot_scene_msgs/SceneObjectArray (BEST_EFFORT)

Publishes
---------
  /saltybot/person_tracks        saltybot_scene_msgs/PersonTrackArray  at publish_hz

Parameters
----------
camera_id            str    "front"   Label stored in every PersonTrack
similarity_threshold float  0.75      Cosine similarity threshold for re-id
stale_timeout_s      float  30.0      Seconds before inactive track is stale
max_tracks           int    20        Maximum concurrent active tracks
publish_hz           float  5.0       Track list publication rate (Hz)
"""

from __future__ import annotations

import time

import cv2
import cv_bridge
import numpy as np

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

from builtin_interfaces.msg import Time as TimeMsg
from sensor_msgs.msg import Image
from saltybot_scene_msgs.msg import (
    SceneObjectArray,
    PersonTrack,
    PersonTrackArray,
)

from ._person_reid import TrackGallery, extract_hsv_histogram


# COCO class 0 = person
_PERSON_CLASS_ID = 0

_BEST_EFFORT_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
)


def _float_to_ros_time(wall_secs: float) -> TimeMsg:
    """Convert a time.time() float to builtin_interfaces/Time."""
    sec_int  = int(wall_secs)
    nanosec  = int((wall_secs - sec_int) * 1e9)
    msg         = TimeMsg()
    msg.sec     = sec_int
    msg.nanosec = nanosec
    return msg


class PersonReidNode(Node):

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

        self.declare_parameter('camera_id',            'front')
        self.declare_parameter('similarity_threshold', 0.75)
        self.declare_parameter('stale_timeout_s',      30.0)
        self.declare_parameter('max_tracks',           20)
        self.declare_parameter('publish_hz',           5.0)

        self._camera_id  = str(self.get_parameter('camera_id').value)
        self._threshold  = float(self.get_parameter('similarity_threshold').value)
        self._timeout    = float(self.get_parameter('stale_timeout_s').value)
        self._max_tracks = int(self.get_parameter('max_tracks').value)
        publish_hz       = float(self.get_parameter('publish_hz').value)

        self._bridge: cv_bridge.CvBridge = cv_bridge.CvBridge()
        self._gallery: TrackGallery      = TrackGallery()
        self._latest_image: np.ndarray | None = None

        self._img_sub = self.create_subscription(
            Image,
            '/camera/color/image_raw',
            self._on_image,
            _BEST_EFFORT_QOS,
        )
        self._det_sub = self.create_subscription(
            SceneObjectArray,
            '/saltybot/scene/objects',
            self._on_detections,
            _BEST_EFFORT_QOS,
        )
        self._pub = self.create_publisher(
            PersonTrackArray,
            '/saltybot/person_tracks',
            10,
        )

        period = 1.0 / max(publish_hz, 0.1)
        self._timer = self.create_timer(period, self._on_timer)

        self.get_logger().info(
            f'person_reid_node ready — '
            f'camera_id={self._camera_id} '
            f'threshold={self._threshold} '
            f'timeout={self._timeout}s '
            f'max_tracks={self._max_tracks}'
        )

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

    def _on_image(self, msg: Image) -> None:
        try:
            self._latest_image = self._bridge.imgmsg_to_cv2(msg, 'bgr8')
        except Exception as exc:
            self.get_logger().warning(f'cv_bridge error: {exc}', throttle_duration_sec=5.0)

    # ── Detection callback ────────────────────────────────────────────────────

    def _on_detections(self, msg: SceneObjectArray) -> None:
        if self._latest_image is None:
            return

        now = time.time()
        img = self._latest_image
        h, w = img.shape[:2]

        for obj in msg.objects:
            if obj.class_id != _PERSON_CLASS_ID:
                continue

            # Crop the person ROI from the cached image
            cx  = obj.bbox.center.position.x
            cy  = obj.bbox.center.position.y
            hsw = obj.bbox.size_x / 2.0
            hsh = obj.bbox.size_y / 2.0
            x1 = max(0, int(cx - hsw))
            y1 = max(0, int(cy - hsh))
            x2 = min(w, int(cx + hsw))
            y2 = min(h, int(cy + hsh))
            if x2 <= x1 or y2 <= y1:
                continue

            crop    = img[y1:y2, x1:x2]
            feature = extract_hsv_histogram(crop)
            bbox    = (cx, cy, obj.bbox.size_x, obj.bbox.size_y)

            track_id, sim = self._gallery.match(feature, self._threshold)

            if track_id is not None:
                self._gallery.update(
                    track_id, self._camera_id, feature, bbox, float(sim), now
                )
            elif len(self._gallery.tracks) < self._max_tracks:
                self._gallery.add(self._camera_id, feature, bbox, now)

    # ── Timer callback — publish + housekeeping ───────────────────────────────

    def _on_timer(self) -> None:
        now = time.time()
        self._gallery.mark_stale(now, self._timeout)
        self._gallery.prune_stale()

        stamp = self.get_clock().now().to_msg()
        out   = PersonTrackArray()
        out.header.stamp    = stamp
        out.header.frame_id = self._camera_id

        for entry in self._gallery.tracks.values():
            t                  = PersonTrack()
            t.header.stamp     = stamp
            t.header.frame_id  = self._camera_id
            t.track_id         = entry.track_id
            t.camera_id        = entry.camera_id
            t.confidence       = float(entry.confidence)
            t.first_seen       = _float_to_ros_time(entry.first_seen)
            t.last_seen        = _float_to_ros_time(entry.last_seen)
            t.is_stale         = entry.is_stale
            # Populate last-known bounding box
            cx, cy, sx, sy     = entry.bbox
            t.bbox.center.position.x = cx
            t.bbox.center.position.y = cy
            t.bbox.size_x            = sx
            t.bbox.size_y            = sy
            out.tracks.append(t)

        self._pub.publish(out)


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


if __name__ == '__main__':
    main()
