"""
aruco_detect_node.py — ArUco marker detection for docking alignment (Issue #627).

Detects all DICT_4X4_50 ArUco markers visible in the RealSense D435i RGB stream,
estimates their 6-DOF poses relative to the camera using estimatePoseSingleMarkers,
and publishes:

  /saltybot/aruco/markers      geometry_msgs/PoseArray        — all detected markers
  /saltybot/aruco/dock_target  geometry_msgs/PoseStamped      — closest dock marker
  /saltybot/aruco/viz          visualization_msgs/MarkerArray — RViz axes overlays
  /saltybot/aruco/status       std_msgs/String                — JSON summary (10 Hz)

Coordinate frame
────────────────
All poses are expressed in camera_color_optical_frame (ROS optical convention):
  +X = right, +Y = down, +Z = forward (into scene)

The dock_target pose gives the docking controller everything it needs:
  position.z  — forward distance to dock (m)      → throttle target
  position.x  — lateral offset (m)                → steer target
  orientation — marker orientation for final align → yaw servo

Dock target selection
─────────────────────
  dock_marker_ids   list[int]  []   — empty = accept any detected marker
                                       non-empty = only these IDs are candidates
  closest_wins      bool       true — among candidates, prefer the nearest
  max_dock_range_m  float      3.0  — ignore markers beyond this distance

estimatePoseSingleMarkers API
──────────────────────────────
  Uses cv2.aruco.estimatePoseSingleMarkers (legacy OpenCV API still present in
  4.x with a deprecation notice).  Falls back to per-marker cv2.solvePnP with
  SOLVEPNP_IPPE_SQUARE if the legacy function is unavailable.
  Both paths use the same camera_matrix / dist_coeffs from CameraInfo.

Subscribes
──────────
  /camera/color/image_raw       sensor_msgs/Image      30 Hz (BGR8)
  /camera/color/camera_info     sensor_msgs/CameraInfo (latched, once)

Publishes
─────────
  /saltybot/aruco/markers       geometry_msgs/PoseArray
  /saltybot/aruco/dock_target   geometry_msgs/PoseStamped
  /saltybot/aruco/viz           visualization_msgs/MarkerArray
  /saltybot/aruco/status        std_msgs/String   (JSON, 10 Hz)

Parameters (see config/aruco_detect_params.yaml)
────────────────────────────────────────────────
  marker_size_m      float   0.10     physical side length of printed markers (m)
  dock_marker_ids    int[]   []       IDs to accept as dock targets (empty = all)
  max_dock_range_m   float   3.0      ignore candidates beyond this (m)
  aruco_dict         str     DICT_4X4_50
  corner_refinement  str     CORNER_REFINE_SUBPIX
  camera_frame       str     camera_color_optical_frame
  draw_debug_image   bool    false    publish annotated image for debugging
  debug_image_topic  str     /saltybot/aruco/debug_image
"""

from __future__ import annotations

import json
import math

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

import numpy as np
import cv2
from cv_bridge import CvBridge

from geometry_msgs.msg import (
    Pose, PoseArray, PoseStamped, Point, Quaternion, Vector3, TransformStamped
)
from sensor_msgs.msg import Image, CameraInfo
from std_msgs.msg import Header, String, ColorRGBA
from visualization_msgs.msg import Marker, MarkerArray

from .aruco_math import MarkerPose, rvec_to_quat, tvec_distance


# ── QoS ───────────────────────────────────────────────────────────────────────
_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)
_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)

# ArUco axis length for viz (fraction of marker size)
_AXIS_LEN_FRAC = 0.5


class ArucoDetectNode(Node):
    """ArUco marker detection + pose estimation node — see module docstring."""

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter('marker_size_m',     0.10)
        self.declare_parameter('dock_marker_ids',   [])       # empty list = all
        self.declare_parameter('max_dock_range_m',  3.0)
        self.declare_parameter('aruco_dict',        'DICT_4X4_50')
        self.declare_parameter('corner_refinement', 'CORNER_REFINE_SUBPIX')
        self.declare_parameter('camera_frame',      'camera_color_optical_frame')
        self.declare_parameter('draw_debug_image',  False)
        self.declare_parameter('debug_image_topic', '/saltybot/aruco/debug_image')

        self._marker_size  = self.get_parameter('marker_size_m').value
        self._dock_ids     = set(self.get_parameter('dock_marker_ids').value)
        self._max_range    = self.get_parameter('max_dock_range_m').value
        self._cam_frame    = self.get_parameter('camera_frame').value
        self._draw_debug   = self.get_parameter('draw_debug_image').value
        self._debug_topic  = self.get_parameter('debug_image_topic').value

        aruco_dict_name    = self.get_parameter('aruco_dict').value
        refine_name        = self.get_parameter('corner_refinement').value

        # ── ArUco detector setup ──────────────────────────────────────────────
        dict_id    = getattr(cv2.aruco, aruco_dict_name, cv2.aruco.DICT_4X4_50)
        params     = cv2.aruco.DetectorParameters()
        # Apply corner refinement
        refine_id  = getattr(cv2.aruco, refine_name, cv2.aruco.CORNER_REFINE_SUBPIX)
        params.cornerRefinementMethod = refine_id

        aruco_dict     = cv2.aruco.getPredefinedDictionary(dict_id)
        self._detector = cv2.aruco.ArucoDetector(aruco_dict, params)

        # Object points for solvePnP fallback (counter-clockwise from top-left)
        half = self._marker_size / 2.0
        self._obj_pts = np.array([
            [-half,  half, 0.0],
            [ half,  half, 0.0],
            [ half, -half, 0.0],
            [-half, -half, 0.0],
        ], dtype=np.float64)

        # ── Camera intrinsics (set on first CameraInfo) ───────────────────────
        self._K:    np.ndarray | None = None
        self._dist: np.ndarray | None = None
        self._bridge = CvBridge()

        # ── Publishers ─────────────────────────────────────────────────────────
        self._pose_array_pub = self.create_publisher(
            PoseArray, '/saltybot/aruco/markers', 10
        )
        self._dock_pub = self.create_publisher(
            PoseStamped, '/saltybot/aruco/dock_target', 10
        )
        self._viz_pub = self.create_publisher(
            MarkerArray, '/saltybot/aruco/viz', 10
        )
        self._status_pub = self.create_publisher(
            String, '/saltybot/aruco/status', 10
        )
        if self._draw_debug:
            self._debug_pub = self.create_publisher(
                Image, self._debug_topic, 5
            )
        else:
            self._debug_pub = None

        # ── Subscriptions ──────────────────────────────────────────────────────
        self.create_subscription(
            CameraInfo,
            '/camera/color/camera_info',
            self._on_camera_info,
            _LATCHED_QOS,
        )
        self.create_subscription(
            Image,
            '/camera/color/image_raw',
            self._on_image,
            _SENSOR_QOS,
        )

        # ── Status timer (10 Hz) ────────────────────────────────────────────────
        self._last_detections: list[MarkerPose] = []
        self.create_timer(0.1, self._on_status_timer)

        self.get_logger().info(
            f'aruco_detect ready — '
            f'dict={aruco_dict_name} '
            f'marker_size={self._marker_size * 100:.0f}cm '
            f'dock_ids={list(self._dock_ids) if self._dock_ids else "any"} '
            f'max_range={self._max_range}m'
        )

    # ── Camera info ───────────────────────────────────────────────────────────

    def _on_camera_info(self, msg: CameraInfo) -> None:
        if self._K is not None:
            return
        self._K    = np.array(msg.k, dtype=np.float64).reshape(3, 3)
        self._dist = np.array(msg.d, dtype=np.float64).reshape(1, -1)
        self.get_logger().info(
            f'camera_info received — '
            f'fx={self._K[0,0]:.1f} fy={self._K[1,1]:.1f}'
        )

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

    def _on_image(self, msg: Image) -> None:
        if self._K is None:
            return

        # ── Decode ────────────────────────────────────────────────────────────
        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as exc:
            self.get_logger().warning(f'image decode error: {exc}',
                                      throttle_duration_sec=5.0)
            return

        # Convert to greyscale for detection (faster, same accuracy)
        gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)

        stamp  = msg.header.stamp

        # ── Detect markers ────────────────────────────────────────────────────
        corners, ids, _ = self._detector.detectMarkers(gray)

        if ids is None or len(ids) == 0:
            self._last_detections = []
            self._publish_empty(stamp)
            return

        ids_flat = ids.flatten().tolist()

        # ── Pose estimation ───────────────────────────────────────────────────
        detections = self._estimate_poses(corners, ids_flat)

        # Filter by range
        detections = [d for d in detections if d.distance_m <= self._max_range]

        self._last_detections = detections

        # ── Publish ───────────────────────────────────────────────────────────
        self._publish_pose_array(detections, stamp)
        self._publish_dock_target(detections, stamp)
        self._publish_viz(detections, stamp)

        # Debug image
        if self._debug_pub is not None:
            self._publish_debug(bgr, corners, ids, detections, stamp)

    # ── Pose estimation ───────────────────────────────────────────────────────

    def _estimate_poses(
        self,
        corners: list,
        ids:     list[int],
    ) -> list[MarkerPose]:
        """
        Estimate 6-DOF pose for each detected marker.

        Uses estimatePoseSingleMarkers (legacy API, still present in cv2 4.x).
        Falls back to per-marker solvePnP(IPPE_SQUARE) if unavailable.
        """
        results: list[MarkerPose] = []

        # ── Try legacy estimatePoseSingleMarkers ──────────────────────────────
        if hasattr(cv2.aruco, 'estimatePoseSingleMarkers'):
            try:
                rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(
                    corners, self._marker_size, self._K, self._dist
                )
                for i, mid in enumerate(ids):
                    rvec = rvecs[i].ravel().astype(np.float64)
                    tvec = tvecs[i].ravel().astype(np.float64)
                    results.append(MarkerPose(
                        marker_id = int(mid),
                        tvec      = tvec,
                        rvec      = rvec,
                        corners   = corners[i].reshape(4, 2).astype(np.float64),
                    ))
                return results
            except Exception as exc:
                self.get_logger().debug(
                    f'estimatePoseSingleMarkers failed ({exc}), '
                    'falling back to solvePnP'
                )

        # ── Fallback: per-marker solvePnP ─────────────────────────────────────
        for i, mid in enumerate(ids):
            img_pts = corners[i].reshape(4, 2).astype(np.float64)
            ok, rvec, tvec = cv2.solvePnP(
                self._obj_pts, img_pts,
                self._K, self._dist,
                flags=cv2.SOLVEPNP_IPPE_SQUARE,
            )
            if not ok:
                continue
            results.append(MarkerPose(
                marker_id = int(mid),
                tvec      = tvec.ravel().astype(np.float64),
                rvec      = rvec.ravel().astype(np.float64),
                corners   = img_pts,
            ))
        return results

    # ── Dock target selection ─────────────────────────────────────────────────

    def _select_dock_target(
        self,
        detections: list[MarkerPose],
    ) -> MarkerPose | None:
        """Select the closest marker that matches dock_marker_ids filter."""
        candidates = [
            d for d in detections
            if not self._dock_ids or d.marker_id in self._dock_ids
        ]
        if not candidates:
            return None
        return min(candidates, key=lambda d: d.distance_m)

    # ── Publishers ────────────────────────────────────────────────────────────

    def _publish_pose_array(
        self,
        detections: list[MarkerPose],
        stamp,
    ) -> None:
        pa              = PoseArray()
        pa.header.stamp    = stamp
        pa.header.frame_id = self._cam_frame

        for d in detections:
            qx, qy, qz, qw = d.quat
            pose               = Pose()
            pose.position      = Point(x=float(d.tvec[0]),
                                       y=float(d.tvec[1]),
                                       z=float(d.tvec[2]))
            pose.orientation   = Quaternion(x=qx, y=qy, z=qz, w=qw)
            pa.poses.append(pose)

        self._pose_array_pub.publish(pa)

    def _publish_dock_target(
        self,
        detections: list[MarkerPose],
        stamp,
    ) -> None:
        target = self._select_dock_target(detections)
        if target is None:
            return

        qx, qy, qz, qw = target.quat
        ps                  = PoseStamped()
        ps.header.stamp     = stamp
        ps.header.frame_id  = self._cam_frame
        ps.pose.position    = Point(x=float(target.tvec[0]),
                                    y=float(target.tvec[1]),
                                    z=float(target.tvec[2]))
        ps.pose.orientation = Quaternion(x=qx, y=qy, z=qz, w=qw)
        self._dock_pub.publish(ps)

    def _publish_viz(
        self,
        detections: list[MarkerPose],
        stamp,
    ) -> None:
        """Publish coordinate-axes MarkerArray for each detected marker."""
        ma      = MarkerArray()
        lifetime = _ros_duration(0.2)

        # Delete-all
        dm              = Marker()
        dm.action       = Marker.DELETEALL
        dm.header.stamp    = stamp
        dm.header.frame_id = self._cam_frame
        ma.markers.append(dm)

        target = self._select_dock_target(detections)
        axis_len = self._marker_size * _AXIS_LEN_FRAC

        for idx, d in enumerate(detections):
            is_target = (target is not None and d.marker_id == target.marker_id)
            cx, cy, cz = float(d.tvec[0]), float(d.tvec[1]), float(d.tvec[2])
            qx, qy, qz, qw = d.quat
            base_id = idx * 10

            # Sphere at marker centre
            sph                       = Marker()
            sph.header.stamp          = stamp
            sph.header.frame_id       = self._cam_frame
            sph.ns                    = 'aruco_centers'
            sph.id                    = base_id
            sph.type                  = Marker.SPHERE
            sph.action                = Marker.ADD
            sph.pose.position         = Point(x=cx, y=cy, z=cz)
            sph.pose.orientation      = Quaternion(x=qx, y=qy, z=qz, w=qw)
            r_s = 0.015 if not is_target else 0.025
            sph.scale                 = Vector3(x=r_s * 2, y=r_s * 2, z=r_s * 2)
            sph.color = (
                ColorRGBA(r=1.0, g=0.2, b=0.2, a=1.0) if is_target else
                ColorRGBA(r=0.2, g=0.8, b=1.0, a=0.9)
            )
            sph.lifetime              = lifetime
            ma.markers.append(sph)

            # Text label: ID + distance
            txt                       = Marker()
            txt.header.stamp          = stamp
            txt.header.frame_id       = self._cam_frame
            txt.ns                    = 'aruco_labels'
            txt.id                    = base_id + 1
            txt.type                  = Marker.TEXT_VIEW_FACING
            txt.action                = Marker.ADD
            txt.pose.position         = Point(x=cx, y=cy - 0.08, z=cz)
            txt.pose.orientation.w    = 1.0
            txt.scale.z               = 0.06
            txt.color                 = ColorRGBA(r=1.0, g=1.0, b=0.0, a=1.0)
            txt.text = f'ID={d.marker_id}\n{d.distance_m:.2f}m'
            if is_target:
                txt.text += '\n[DOCK]'
            txt.lifetime              = lifetime
            ma.markers.append(txt)

        self._viz_pub.publish(ma)

    def _publish_empty(self, stamp) -> None:
        pa              = PoseArray()
        pa.header.stamp    = stamp
        pa.header.frame_id = self._cam_frame
        self._pose_array_pub.publish(pa)
        self._publish_viz([], stamp)

    def _publish_debug(self, bgr, corners, ids, detections, stamp) -> None:
        """Annotate and publish debug image with detected markers drawn."""
        debug = bgr.copy()
        if ids is not None and len(ids) > 0:
            cv2.aruco.drawDetectedMarkers(debug, corners, ids)

        # Draw axes for each detection using the estimated poses
        for d in detections:
            try:
                cv2.drawFrameAxes(
                    debug, self._K, self._dist,
                    d.rvec.reshape(3, 1),
                    d.tvec.reshape(3, 1),
                    self._marker_size * _AXIS_LEN_FRAC,
                )
            except Exception:
                pass

        try:
            img_msg = self._bridge.cv2_to_imgmsg(debug, encoding='bgr8')
            img_msg.header.stamp    = stamp
            img_msg.header.frame_id = self._cam_frame
            self._debug_pub.publish(img_msg)
        except Exception:
            pass

    # ── Status timer ─────────────────────────────────────────────────────────

    def _on_status_timer(self) -> None:
        detections = self._last_detections
        target     = self._select_dock_target(detections)

        markers_info = [
            {
                'id':         d.marker_id,
                'distance_m': round(d.distance_m,  3),
                'yaw_deg':    round(math.degrees(d.yaw_rad), 2),
                'lateral_m':  round(d.lateral_m,   3),
                'forward_m':  round(d.forward_m,   3),
                'is_target':  target is not None and d.marker_id == target.marker_id,
            }
            for d in detections
        ]

        status = {
            'detected_count':  len(detections),
            'dock_target_id':  target.marker_id  if target else None,
            'dock_distance_m': round(target.distance_m, 3) if target else None,
            'dock_yaw_deg':    round(math.degrees(target.yaw_rad), 2) if target else None,
            'dock_lateral_m':  round(target.lateral_m, 3) if target else None,
            'markers':         markers_info,
        }
        msg      = String()
        msg.data = json.dumps(status)
        self._status_pub.publish(msg)


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

def _ros_duration(seconds: float):
    from rclpy.duration import Duration
    sec  = int(seconds)
    nsec = int((seconds - sec) * 1e9)
    return Duration(seconds=sec, nanoseconds=nsec).to_msg()


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


if __name__ == '__main__':
    main()
