"""
apriltag_node.py — AprilTag (DICT_APRILTAG_36h11) landmark detector.

Detects AprilTags in CSI camera frames and publishes results at 10 Hz.
When camera_info is available, estimates full 6-DOF pose per tag via
cv2.solvePnP (SOLVEPNP_IPPE_SQUARE).

Supports OpenCV 4.6 (legacy detectMarkers) and 4.7+ (ArucoDetector).

Subscribes:
  <image_topic>                      sensor_msgs/Image      default /camera/color/image_raw
  <image_topic_dir>/camera_info      sensor_msgs/CameraInfo  latched intrinsics

Publishes:
  /saltybot/apriltags                saltybot_apriltag_msgs/DetectedTagArray  10 Hz

Parameters:
  image_topic     str    '/camera/color/image_raw'
  publish_hz      float  10.0
  tag_size_m      float  0.10   physical tag side in metres (for pose estimation)
  min_margin      float  10.0   minimum corner-span pixels to keep a detection
"""

from __future__ import annotations

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 sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import Point, Pose

try:
    from saltybot_apriltag_msgs.msg import DetectedTag, DetectedTagArray
    _MSGS_OK = True
except ImportError:
    _MSGS_OK = False

from ._aruco_utils import ArucoBackend, rvec_tvec_to_pose


_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,
)

_FAMILY = 'DICT_APRILTAG_36h11'


class AprilTagNode(Node):

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

        self.declare_parameter('image_topic', '/camera/color/image_raw')
        self.declare_parameter('publish_hz',   10.0)
        self.declare_parameter('tag_size_m',   0.10)
        self.declare_parameter('min_margin',   10.0)

        img_topic        = self.get_parameter('image_topic').value
        pub_hz           = self.get_parameter('publish_hz').value
        self._tag_size   = self.get_parameter('tag_size_m').value
        self._min_margin = self.get_parameter('min_margin').value

        self._bridge  = CvBridge()
        self._backend = ArucoBackend(self.get_logger())

        self._K:    np.ndarray | None = None
        self._dist: np.ndarray | None = None
        self._latest_img: Image | None = None

        # Derive camera_info topic from image topic namespace
        info_topic = img_topic.rsplit('/', 1)[0] + '/camera_info'

        self.create_subscription(Image,      img_topic,   self._on_image,       _SENSOR_QOS)
        self.create_subscription(CameraInfo, info_topic,  self._on_camera_info, _LATCHED_QOS)

        if _MSGS_OK:
            self._pub = self.create_publisher(
                DetectedTagArray, '/saltybot/apriltags', 10
            )
        else:
            self._pub = None
            self.get_logger().warning(
                '[apriltag] saltybot_apriltag_msgs not built — '
                'DetectedTagArray not published'
            )

        self.create_timer(1.0 / pub_hz, self._tick)

        self.get_logger().info(
            f'apriltag_detector ready — dict={_FAMILY} '
            f'tag_size={self._tag_size}m topic={img_topic} hz={pub_hz}'
        )

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

    def _on_image(self, msg: Image) -> None:
        self._latest_img = msg

    def _on_camera_info(self, msg: CameraInfo) -> None:
        self._K    = np.array(msg.k, dtype=np.float64).reshape(3, 3)
        self._dist = np.array(msg.d, dtype=np.float64)

    # ── 10 Hz tick ────────────────────────────────────────────────────────────

    def _tick(self) -> None:
        if self._latest_img is None or self._pub is None:
            return

        try:
            gray = self._bridge.imgmsg_to_cv2(
                self._latest_img, desired_encoding='mono8'
            )
        except Exception as exc:
            self.get_logger().warning(f'[apriltag] decode error: {exc}')
            return

        corners, ids = self._backend.detect(gray)

        stamp = self.get_clock().now().to_msg()
        arr   = self._build_array(corners, ids, self._latest_img.header, stamp)
        self._pub.publish(arr)

    # ── Message builder ───────────────────────────────────────────────────────

    def _build_array(self, corners, ids, src_header, stamp) -> 'DetectedTagArray':
        arr = DetectedTagArray()
        arr.header.stamp    = stamp
        arr.header.frame_id = src_header.frame_id

        if ids is None or len(ids) == 0:
            arr.count = 0
            return arr

        half    = self._tag_size / 2.0
        obj_pts = np.array([
            [-half,  half, 0],
            [ half,  half, 0],
            [ half, -half, 0],
            [-half, -half, 0],
        ], dtype=np.float64)

        for corner_set, tag_id in zip(corners, ids.flatten()):
            pts    = corner_set.reshape(4, 2)
            margin = float(np.mean([
                np.linalg.norm(pts[i] - pts[(i + 1) % 4]) for i in range(4)
            ]))
            if margin < self._min_margin:
                continue

            cx = float(np.mean(pts[:, 0]))
            cy = float(np.mean(pts[:, 1]))

            tag = DetectedTag()
            tag.header          = arr.header
            tag.tag_id          = int(tag_id)
            tag.family          = _FAMILY
            tag.decision_margin = margin
            tag.corners         = pts.flatten().tolist()
            tag.center          = Point(x=cx, y=cy, z=0.0)

            if self._K is not None and self._dist is not None:
                ok, rvec, tvec = cv2.solvePnP(
                    obj_pts, pts.astype(np.float64),
                    self._K, self._dist,
                    flags=cv2.SOLVEPNP_IPPE_SQUARE,
                )
                if ok:
                    tx, ty, tz, qx, qy, qz, qw = rvec_tvec_to_pose(rvec, tvec)
                    tag.pose.position.x    = tx
                    tag.pose.position.y    = ty
                    tag.pose.position.z    = tz
                    tag.pose.orientation.x = qx
                    tag.pose.orientation.y = qy
                    tag.pose.orientation.z = qz
                    tag.pose.orientation.w = qw
                    tag.pose_valid         = True
                else:
                    tag.pose_valid = False
            else:
                tag.pose_valid = False

            arr.tags.append(tag)

        arr.count = len(arr.tags)
        return arr


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


if __name__ == '__main__':
    main()
