"""
_color_segmenter.py — HSV color segmentation helpers (no ROS2 deps).

Algorithm
---------
For each requested color:
  1. Convert BGR → HSV (OpenCV: H∈[0,180], S∈[0,255], V∈[0,255])
  2. Build a binary mask via cv2.inRange using the color's HSV bounds.
     Red wraps around H=0/180 so two ranges are OR-combined.
  3. Morphological open (3×3) to remove noise.
  4. Find external contours; filter by min_area_px.
  5. Return ColorBlob NamedTuples — one per surviving contour.

confidence is the contour area divided by the bounding-rectangle area
(how "filled" the bounding box is), clamped to [0, 1].

Public API
----------
  HsvRange(h_lo, h_hi, s_lo, s_hi, v_lo, v_hi)
  ColorBlob(color_name, confidence, cx, cy, w, h, area_px, contour_id)
  COLOR_RANGES : Dict[str, List[HsvRange]]   — default per-color HSV ranges
  mask_for_color(hsv, color_name) -> np.ndarray   — uint8 binary mask
  find_color_blobs(bgr, active_colors, min_area_px, max_blobs_per_color) -> List[ColorBlob]
"""

from __future__ import annotations

from typing import Dict, List, NamedTuple

import numpy as np


# ── Data types ────────────────────────────────────────────────────────────────

class HsvRange(NamedTuple):
    """Single HSV band (OpenCV: H∈[0,180], S/V∈[0,255])."""
    h_lo: int
    h_hi: int
    s_lo: int
    s_hi: int
    v_lo: int
    v_hi: int


class ColorBlob(NamedTuple):
    """One detected color object in image coordinates."""
    color_name:  str
    confidence:  float   # contour_area / bbox_area  (0–1)
    cx:          float   # bbox centre x (pixels)
    cy:          float   # bbox centre y (pixels)
    w:           float   # bbox width  (pixels)
    h:           float   # bbox height (pixels)
    area_px:     float   # contour area (pixels²)
    contour_id:  int     # 0-based index within this color in this frame


# ── Default per-color HSV ranges ──────────────────────────────────────────────
# Two ranges are used for red (wraps at 0/180).
# S_lo=60, V_lo=50 to ignore desaturated / near-black pixels.

COLOR_RANGES: Dict[str, List[HsvRange]] = {
    'red': [
        HsvRange(h_lo=0,   h_hi=10,  s_lo=60, s_hi=255, v_lo=50, v_hi=255),
        HsvRange(h_lo=170, h_hi=180, s_lo=60, s_hi=255, v_lo=50, v_hi=255),
    ],
    'green': [
        HsvRange(h_lo=35, h_hi=85,  s_lo=60, s_hi=255, v_lo=50, v_hi=255),
    ],
    'blue': [
        HsvRange(h_lo=90, h_hi=130, s_lo=60, s_hi=255, v_lo=50, v_hi=255),
    ],
    'yellow': [
        HsvRange(h_lo=18, h_hi=38,  s_lo=60, s_hi=255, v_lo=80, v_hi=255),
    ],
    'orange': [
        HsvRange(h_lo=8,  h_hi=20,  s_lo=80, s_hi=255, v_lo=80, v_hi=255),
    ],
}

# Structuring element for morphological open (noise removal)
_MORPH_KERNEL = None


def _get_morph_kernel():
    import cv2
    global _MORPH_KERNEL
    if _MORPH_KERNEL is None:
        _MORPH_KERNEL = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    return _MORPH_KERNEL


# ── Public helpers ─────────────────────────────────────────────────────────────

def mask_for_color(hsv: np.ndarray, color_name: str) -> np.ndarray:
    """
    Return a uint8 binary mask (255=foreground) for *color_name* in the HSV image.

    Parameters
    ----------
    hsv        : (H, W, 3) uint8 ndarray in OpenCV HSV format (H∈[0,180])
    color_name : one of COLOR_RANGES keys

    Returns
    -------
    (H, W) uint8 ndarray
    """
    import cv2

    ranges = COLOR_RANGES.get(color_name)
    if not ranges:
        raise ValueError(f'Unknown color: {color_name!r}. Known: {list(COLOR_RANGES)}')

    mask = np.zeros(hsv.shape[:2], dtype=np.uint8)
    for r in ranges:
        lo = np.array([r.h_lo, r.s_lo, r.v_lo], dtype=np.uint8)
        hi = np.array([r.h_hi, r.s_hi, r.v_hi], dtype=np.uint8)
        mask |= cv2.inRange(hsv, lo, hi)

    return cv2.morphologyEx(mask, cv2.MORPH_OPEN, _get_morph_kernel())


def find_color_blobs(
    bgr:                 np.ndarray,
    active_colors:       List[str] | None = None,
    min_area_px:         float = 200.0,
    max_blobs_per_color: int   = 10,
) -> List[ColorBlob]:
    """
    Detect HSV-segmented color blobs in a BGR image.

    Parameters
    ----------
    bgr                 : (H, W, 3) uint8 BGR ndarray
    active_colors       : color names to detect; None → all COLOR_RANGES keys
    min_area_px         : minimum contour area to report (pixels²)
    max_blobs_per_color : keep at most this many blobs per color (largest first)

    Returns
    -------
    List[ColorBlob] — may be empty; contour_id is 0-based within each color
    """
    import cv2

    if active_colors is None:
        active_colors = list(COLOR_RANGES.keys())

    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    blobs: List[ColorBlob] = []

    for color_name in active_colors:
        mask = mask_for_color(hsv, color_name)
        contours, _ = cv2.findContours(
            mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Sort largest first so max_blobs_per_color keeps the significant ones
        contours = sorted(contours, key=cv2.contourArea, reverse=True)

        blob_idx = 0
        for cnt in contours:
            if blob_idx >= max_blobs_per_color:
                break

            area = cv2.contourArea(cnt)
            if area < min_area_px:
                break   # already sorted, no need to continue

            x, y, bw, bh = cv2.boundingRect(cnt)
            bbox_area = float(bw * bh)
            confidence = float(area / bbox_area) if bbox_area > 0 else 0.0
            confidence = min(1.0, max(0.0, confidence))

            blobs.append(ColorBlob(
                color_name=color_name,
                confidence=confidence,
                cx=float(x + bw / 2.0),
                cy=float(y + bh / 2.0),
                w=float(bw),
                h=float(bh),
                area_px=float(area),
                contour_id=blob_idx,
            ))
            blob_idx += 1

    return blobs
