"""
test_apriltag.py — Unit tests for AprilTag detector helpers.

Runs without ROS2 (no rclpy). Tests cover:
  - _rvec_tvec_to_pose: rotation/translation → Pose quaternion
  - _ArucoBackend: detection on a synthetic rendered tag image
  - AprilTagNode._build_array: message construction from mock detections
"""

from __future__ import annotations

import math
import sys
import os

import numpy as np
import pytest
import cv2

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_apriltag._aruco_utils import rvec_tvec_to_pose as _rvec_tvec_to_pose, ArucoBackend as _ArucoBackend


# ── _rvec_tvec_to_pose ────────────────────────────────────────────────────────

class TestRvecTvecToPose:
    # _rvec_tvec_to_pose returns (tx, ty, tz, qx, qy, qz, qw)

    def test_identity_rotation_gives_unit_quaternion(self):
        tx, ty, tz, qx, qy, qz, qw = _rvec_tvec_to_pose(np.zeros(3), np.zeros(3))
        assert qw == pytest.approx(1.0, abs=1e-6)
        assert qx == pytest.approx(0.0, abs=1e-6)
        assert qy == pytest.approx(0.0, abs=1e-6)
        assert qz == pytest.approx(0.0, abs=1e-6)

    def test_translation_passed_through(self):
        tvec = np.array([1.5, -0.3, 2.0])
        tx, ty, tz, *_ = _rvec_tvec_to_pose(np.zeros(3), tvec)
        assert tx == pytest.approx(1.5)
        assert ty == pytest.approx(-0.3)
        assert tz == pytest.approx(2.0)

    def test_90deg_yaw_quaternion(self):
        # 90° rotation around Z axis → qw=cos(45°), qz=sin(45°)
        rvec = np.array([0.0, 0.0, math.pi / 2])
        tx, ty, tz, qx, qy, qz, qw = _rvec_tvec_to_pose(rvec, np.zeros(3))
        assert qw == pytest.approx(math.cos(math.pi / 4), abs=1e-5)
        assert qz == pytest.approx(math.sin(math.pi / 4), abs=1e-5)
        assert qx == pytest.approx(0.0, abs=1e-5)
        assert qy == pytest.approx(0.0, abs=1e-5)

    def test_quaternion_is_unit_norm(self):
        rvec = np.array([0.3, 0.7, -1.1])
        tvec = np.array([0.5, 0.5, 0.5])
        tx, ty, tz, qx, qy, qz, qw = _rvec_tvec_to_pose(rvec, tvec)
        norm = math.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
        assert norm == pytest.approx(1.0, abs=1e-5)

    def test_180deg_x_rotation(self):
        rvec = np.array([math.pi, 0.0, 0.0])
        tx, ty, tz, qx, qy, qz, qw = _rvec_tvec_to_pose(rvec, np.zeros(3))
        norm = math.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
        assert norm == pytest.approx(1.0, abs=1e-5)


# ── _ArucoBackend ─────────────────────────────────────────────────────────────

class _FakeLogger:
    def info(self, msg): pass
    def warning(self, msg): pass
    def debug(self, msg): pass


def _render_apriltag_36h11(tag_id: int = 0, size: int = 200) -> np.ndarray:
    """
    Render an AprilTag DICT_APRILTAG_36h11 tag into a greyscale image.
    Returns an image large enough for reliable detection.
    """
    dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_APRILTAG_36h11)
    # Draw the marker (OpenCV ≥ 4.6)
    tag_img = np.zeros((size, size), dtype=np.uint8)
    try:
        cv2.aruco.generateImageMarker(dictionary, tag_id, size, tag_img, 1)
    except AttributeError:
        # Older API
        tag_img = cv2.aruco.drawMarker(dictionary, tag_id, size)

    # Embed in white canvas with border
    canvas = np.ones((size + 80, size + 80), dtype=np.uint8) * 255
    canvas[40:40 + size, 40:40 + size] = tag_img
    return canvas


class TestArucoBackend:

    def test_backend_initialises(self):
        backend = _ArucoBackend(_FakeLogger())
        assert backend is not None

    def test_detect_returns_tuple(self):
        backend = _ArucoBackend(_FakeLogger())
        gray = np.ones((480, 640), dtype=np.uint8) * 200
        corners, ids = backend.detect(gray)
        assert isinstance(corners, (list, tuple))

    def test_blank_image_no_detections(self):
        backend = _ArucoBackend(_FakeLogger())
        gray = np.ones((480, 640), dtype=np.uint8) * 200
        _, ids = backend.detect(gray)
        assert ids is None or len(ids) == 0

    def test_rendered_tag_detected(self):
        backend  = _ArucoBackend(_FakeLogger())
        tag_img  = _render_apriltag_36h11(tag_id=5, size=200)
        corners, ids = backend.detect(tag_img)
        assert ids is not None and len(ids) == 1
        assert int(ids.flatten()[0]) == 5

    def test_rendered_tag_corners_shape(self):
        backend = _ArucoBackend(_FakeLogger())
        tag_img = _render_apriltag_36h11(tag_id=0, size=200)
        corners, ids = backend.detect(tag_img)
        assert ids is not None
        assert len(corners) == 1
        assert corners[0].reshape(4, 2).shape == (4, 2)

    def test_multiple_tags_detected(self):
        """Place two tags side-by-side and expect both found."""
        backend = _ArucoBackend(_FakeLogger())
        t0 = _render_apriltag_36h11(tag_id=0, size=150)
        t1 = _render_apriltag_36h11(tag_id=1, size=150)
        canvas = np.ones((t0.shape[0], t0.shape[1] * 2 + 20), dtype=np.uint8) * 255
        canvas[:, :t0.shape[1]] = t0
        canvas[:, t0.shape[1] + 20:] = t1
        _, ids = backend.detect(canvas)
        assert ids is not None and len(ids) == 2
        assert set(ids.flatten().tolist()) == {0, 1}


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
