"""
test_color_segmenter.py — Unit tests for HSV color segmentation helpers (no ROS2 required).

Covers:
  HsvRange / ColorBlob:
    - NamedTuple fields accessible by name
    - confidence clamped to [0,1]

  mask_for_color:
    - pure red image → red mask fully white
    - pure red image → green mask fully black
    - pure green image → green mask fully white
    - pure blue image → blue mask fully white
    - pure yellow image → yellow mask non-empty
    - pure orange image → orange mask non-empty
    - red hue wrap-around detected from both HSV bands
    - unknown color name raises ValueError
    - mask is uint8
    - mask shape matches input

  find_color_blobs — output contract:
    - returns list
    - empty list on blank (no-color) image
    - empty list when min_area_px larger than any contour

  find_color_blobs — detection:
    - large red rectangle detected as red blob
    - large green rectangle detected as green blob
    - large blue rectangle detected as blue blob
    - detected blob color_name matches requested color
    - contour_id is 0 for first blob
    - confidence in [0, 1]
    - cx, cy within image bounds
    - w, h > 0 for detected blob
    - area_px > 0 for detected blob

  find_color_blobs — filtering:
    - active_colors=None detects all colors when present
    - only requested colors returned when active_colors restricted
    - max_blobs_per_color limits output count
    - two separate red blobs both detected when max_blobs=2
    - smaller blob filtered when min_area_px high

  find_color_blobs — multi-color:
    - image with red + green regions → both detected
"""

import sys
import os

import numpy as np
import pytest

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

from saltybot_bringup._color_segmenter import (
    HsvRange,
    ColorBlob,
    COLOR_RANGES,
    mask_for_color,
    find_color_blobs,
)


# ── Image factories ───────────────────────────────────────────────────────────

def _solid_bgr(b, g, r, h=64, w=64) -> np.ndarray:
    """Solid BGR image."""
    img = np.zeros((h, w, 3), dtype=np.uint8)
    img[:, :] = (b, g, r)
    return img


def _blank(h=64, w=64) -> np.ndarray:
    """All-black image (nothing to detect)."""
    return np.zeros((h, w, 3), dtype=np.uint8)


def _image_with_rect(bg_bgr, rect_bgr, rect_slice_r, rect_slice_c, h=128, w=128) -> np.ndarray:
    """Background colour with a filled rectangle."""
    img = np.zeros((h, w, 3), dtype=np.uint8)
    img[:, :] = bg_bgr
    img[rect_slice_r, rect_slice_c] = rect_bgr
    return img


# Canonical solid color BGR values (saturated, in-range for HSV thresholds)
_RED_BGR    = (0,   0,   200)   # BGR pure red
_GREEN_BGR  = (0,   200, 0  )   # BGR pure green
_BLUE_BGR   = (200, 0,   0  )   # BGR pure blue
_YELLOW_BGR = (0,   220, 220)   # BGR yellow
_ORANGE_BGR = (0,   140, 220)   # BGR orange


# ── HsvRange / ColorBlob types ────────────────────────────────────────────────

class TestTypes:

    def test_hsv_range_fields(self):
        r = HsvRange(0, 10, 60, 255, 50, 255)
        assert r.h_lo == 0 and r.h_hi == 10
        assert r.s_lo == 60 and r.s_hi == 255
        assert r.v_lo == 50 and r.v_hi == 255

    def test_color_blob_fields(self):
        b = ColorBlob('red', 0.8, 32.0, 32.0, 20.0, 20.0, 300.0, 0)
        assert b.color_name == 'red'
        assert b.confidence == pytest.approx(0.8)
        assert b.contour_id == 0

    def test_color_ranges_contains_all_defaults(self):
        for color in ('red', 'green', 'blue', 'yellow', 'orange'):
            assert color in COLOR_RANGES
            assert len(COLOR_RANGES[color]) >= 1


# ── mask_for_color ────────────────────────────────────────────────────────────

class TestMaskForColor:

    def test_mask_is_uint8(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_RED_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'red')
        assert m.dtype == np.uint8

    def test_mask_shape_matches_input(self):
        import cv2
        bgr = _solid_bgr(*_RED_BGR, h=48, w=80)
        hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'red')
        assert m.shape == (48, 80)

    def test_pure_red_gives_red_mask_nonzero(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_RED_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'red')
        assert m.any(), 'red mask should be non-empty for red image'

    def test_pure_red_gives_green_mask_empty(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_RED_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'green')
        assert not m.any(), 'green mask should be empty for red image'

    def test_pure_green_gives_green_mask_nonzero(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_GREEN_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'green')
        assert m.any()

    def test_pure_blue_gives_blue_mask_nonzero(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_BLUE_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'blue')
        assert m.any()

    def test_pure_yellow_gives_yellow_mask_nonzero(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_YELLOW_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'yellow')
        assert m.any()

    def test_pure_orange_gives_orange_mask_nonzero(self):
        import cv2
        hsv = cv2.cvtColor(_solid_bgr(*_ORANGE_BGR), cv2.COLOR_BGR2HSV)
        m = mask_for_color(hsv, 'orange')
        assert m.any()

    def test_unknown_color_raises(self):
        import cv2
        hsv = cv2.cvtColor(_blank(), cv2.COLOR_BGR2HSV)
        with pytest.raises(ValueError, match='Unknown color'):
            mask_for_color(hsv, 'purple')

    def test_red_detected_in_high_hue_band(self):
        """A near-180-hue red pixel should still trigger the red mask."""
        import cv2
        # HSV (175, 200, 200) = high-hue red (wrap-around band)
        hsv = np.full((32, 32, 3), (175, 200, 200), dtype=np.uint8)
        m = mask_for_color(hsv, 'red')
        assert m.any(), 'high-hue red not detected'


# ── find_color_blobs — output contract ───────────────────────────────────────

class TestFindColorBlobsContract:

    def test_returns_list(self):
        result = find_color_blobs(_blank())
        assert isinstance(result, list)

    def test_blank_image_returns_empty(self):
        result = find_color_blobs(_blank())
        assert result == []

    def test_min_area_filter_removes_all(self):
        """Request a min area larger than the entire image → no blobs."""
        bgr = _solid_bgr(*_RED_BGR, h=32, w=32)
        result = find_color_blobs(bgr, active_colors=['red'], min_area_px=1e9)
        assert result == []


# ── find_color_blobs — detection ─────────────────────────────────────────────

class TestFindColorBlobsDetection:

    def _large_rect(self, color_bgr, color_name) -> np.ndarray:
        """100×100 image with a 60×60 solid-color rectangle centred."""
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = color_bgr
        return img

    def test_red_rect_detected(self):
        blobs = find_color_blobs(self._large_rect(_RED_BGR, 'red'), active_colors=['red'])
        assert len(blobs) >= 1
        assert blobs[0].color_name == 'red'

    def test_green_rect_detected(self):
        blobs = find_color_blobs(self._large_rect(_GREEN_BGR, 'green'), active_colors=['green'])
        assert len(blobs) >= 1
        assert blobs[0].color_name == 'green'

    def test_blue_rect_detected(self):
        blobs = find_color_blobs(self._large_rect(_BLUE_BGR, 'blue'), active_colors=['blue'])
        assert len(blobs) >= 1
        assert blobs[0].color_name == 'blue'

    def test_first_contour_id_is_zero(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'])
        assert blobs[0].contour_id == 0

    def test_confidence_in_range(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _GREEN_BGR
        blobs = find_color_blobs(img, active_colors=['green'])
        assert blobs
        assert 0.0 <= blobs[0].confidence <= 1.0

    def test_cx_within_image(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _BLUE_BGR
        blobs = find_color_blobs(img, active_colors=['blue'])
        assert blobs
        assert 0.0 <= blobs[0].cx <= 100.0

    def test_cy_within_image(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _BLUE_BGR
        blobs = find_color_blobs(img, active_colors=['blue'])
        assert blobs
        assert 0.0 <= blobs[0].cy <= 100.0

    def test_w_positive(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'])
        assert blobs[0].w > 0

    def test_h_positive(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'])
        assert blobs[0].h > 0

    def test_area_px_positive(self):
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'])
        assert blobs[0].area_px > 0

    def test_area_px_reasonable(self):
        """area_px should be roughly within the rectangle we drew."""
        img = _blank(h=100, w=100)
        img[20:80, 20:80] = _GREEN_BGR   # 60×60 = 3600 px
        blobs = find_color_blobs(img, active_colors=['green'], min_area_px=100.0)
        assert blobs
        assert 1000 <= blobs[0].area_px <= 4000


# ── find_color_blobs — filtering ─────────────────────────────────────────────

class TestFindColorBlobsFiltering:

    def test_active_colors_none_detects_all(self):
        """Image with red+green patches → both found when active_colors=None."""
        img = _blank(h=128, w=128)
        img[10:50, 10:50] = _RED_BGR
        img[10:50, 70:110] = _GREEN_BGR
        blobs = find_color_blobs(img, active_colors=None, min_area_px=100.0)
        names = {b.color_name for b in blobs}
        assert 'red' in names
        assert 'green' in names

    def test_restricted_active_colors(self):
        """Only red requested → no green blobs returned."""
        img = _blank(h=128, w=128)
        img[10:50, 10:50] = _RED_BGR
        img[10:50, 70:110] = _GREEN_BGR
        blobs = find_color_blobs(img, active_colors=['red'], min_area_px=100.0)
        assert all(b.color_name == 'red' for b in blobs)

    def test_max_blobs_per_color_limits(self):
        """Four separate red rectangles but max_blobs=2 → at most 2 blobs."""
        img = _blank(h=200, w=200)
        img[10:40, 10:40]   = _RED_BGR
        img[10:40, 80:110]  = _RED_BGR
        img[100:130, 10:40] = _RED_BGR
        img[100:130, 80:110] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'],
                                 min_area_px=100.0, max_blobs_per_color=2)
        red_blobs = [b for b in blobs if b.color_name == 'red']
        assert len(red_blobs) <= 2

    def test_two_blobs_detected_when_max_allows(self):
        """Two red rectangles detected when max_blobs_per_color >= 2."""
        img = _blank(h=200, w=200)
        img[10:60, 10:60]   = _RED_BGR
        img[10:60, 130:180] = _RED_BGR
        blobs = find_color_blobs(img, active_colors=['red'],
                                 min_area_px=100.0, max_blobs_per_color=10)
        red_blobs = [b for b in blobs if b.color_name == 'red']
        assert len(red_blobs) >= 2

    def test_small_blob_filtered_by_min_area(self):
        """Small 5×5 red patch filtered by min_area_px=500."""
        img = _blank(h=64, w=64)
        img[28:33, 28:33] = _RED_BGR   # 5×5 = 25 px contour area
        blobs = find_color_blobs(img, active_colors=['red'], min_area_px=500.0)
        assert blobs == []


# ── find_color_blobs — multi-color ───────────────────────────────────────────

class TestFindColorBlobsMultiColor:

    def test_red_and_green_in_same_image(self):
        img = _blank(h=128, w=128)
        img[10:60, 10:60]   = _RED_BGR
        img[10:60, 68:118]  = _GREEN_BGR
        blobs = find_color_blobs(img, active_colors=['red', 'green'], min_area_px=100.0)
        names = {b.color_name for b in blobs}
        assert 'red' in names, 'red blob should be detected'
        assert 'green' in names, 'green blob should be detected'

    def test_contour_ids_per_color_start_at_zero(self):
        """contour_id should be 0 for the first (largest) blob of each color."""
        img = _blank(h=200, w=200)
        img[10:80, 10:80]   = _RED_BGR
        img[10:80, 110:180] = _BLUE_BGR
        blobs = find_color_blobs(img, active_colors=['red', 'blue'], min_area_px=100.0)
        for color in ('red', 'blue'):
            first = next((b for b in blobs if b.color_name == color), None)
            assert first is not None, f'{color} blob not found'
            assert first.contour_id == 0, f'{color} first blob contour_id != 0'


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