"""
Color detection for OCR word boxes.

Detects the text color of existing OCR words and recovers colored text
regions (e.g. red markers, blue headings) that standard OCR may have missed.

Standard OCR (Tesseract, PaddleOCR) binarises images before processing,
destroying all color information.  This module adds it back by sampling
HSV pixel values at word-box positions and finding colored regions that
no word-box covers.

Lizenz: Apache 2.0 (kommerziell nutzbar)
DATENSCHUTZ: Alle Verarbeitung erfolgt lokal.
"""

import logging
from typing import Dict, List, Tuple

import cv2
import numpy as np

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# HSV color ranges  (OpenCV:  H 0-180,  S 0-255,  V 0-255)
# ---------------------------------------------------------------------------

_COLOR_RANGES: Dict[str, List[Tuple[np.ndarray, np.ndarray]]] = {
    "red": [
        (np.array([0, 70, 50]), np.array([10, 255, 255])),
        (np.array([170, 70, 50]), np.array([180, 255, 255])),
    ],
    "orange": [
        (np.array([10, 70, 50]), np.array([25, 255, 255])),
    ],
    "yellow": [
        (np.array([25, 70, 50]), np.array([35, 255, 255])),
    ],
    "green": [
        (np.array([35, 70, 50]), np.array([85, 255, 255])),
    ],
    "blue": [
        (np.array([100, 70, 50]), np.array([130, 255, 255])),
    ],
    "purple": [
        (np.array([130, 70, 50]), np.array([170, 255, 255])),
    ],
}

_COLOR_HEX: Dict[str, str] = {
    "black": "#000000",
    "gray": "#6b7280",
    "red": "#dc2626",
    "orange": "#ea580c",
    "yellow": "#ca8a04",
    "green": "#16a34a",
    "blue": "#2563eb",
    "purple": "#9333ea",
}


def _hue_to_color_name(hue: float) -> str:
    """Map OpenCV hue (0-180) to a color name."""
    if hue < 10 or hue > 170:
        return "red"
    if hue < 25:
        return "orange"
    if hue < 35:
        return "yellow"
    if hue < 85:
        return "green"
    if hue < 130:
        return "blue"
    return "purple"


# ---------------------------------------------------------------------------
# 1.  Color annotation for existing word boxes
# ---------------------------------------------------------------------------

def detect_word_colors(
    img_bgr: np.ndarray,
    word_boxes: List[Dict],
    sat_threshold: int = 55,
    min_sat_ratio: float = 0.25,
) -> None:
    """Annotate each word_box in-place with its detected text color.

    Adds ``color`` (hex string) and ``color_name`` (e.g. 'red', 'black')
    keys to each dict.

    Algorithm per word:
      1. Crop the word region from the image.
      2. Otsu-threshold for text/background separation.
      3. Sample background color from border pixels of the crop.
      4. Remove text pixels that match the background (avoids colored
         backgrounds like blue boxes leaking into the result).
      5. Use **median** hue (robust to outliers) and require a minimum
         ratio of saturated pixels before classifying as colored.
    """
    if img_bgr is None or not word_boxes:
        return

    img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
    img_h, img_w = img_bgr.shape[:2]

    colored_count = 0

    for wb in word_boxes:
        x1 = max(0, int(wb["left"]))
        y1 = max(0, int(wb["top"]))
        x2 = min(img_w, int(wb["left"] + wb["width"]))
        y2 = min(img_h, int(wb["top"] + wb["height"]))

        if x2 <= x1 or y2 <= y1:
            wb["color"] = _COLOR_HEX["black"]
            wb["color_name"] = "black"
            continue

        crop_hsv = img_hsv[y1:y2, x1:x2]
        crop_bgr = img_bgr[y1:y2, x1:x2]
        crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
        ch, cw = crop_hsv.shape[:2]

        # --- Text mask: Otsu (adaptive) + high-saturation pixels ---
        _, dark_mask = cv2.threshold(
            crop_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU,
        )
        sat_mask = (crop_hsv[:, :, 1] > sat_threshold).astype(np.uint8) * 255
        text_mask = cv2.bitwise_or(dark_mask, sat_mask)

        text_pixels = crop_hsv[text_mask > 0]

        if len(text_pixels) < 3:
            wb["color"] = _COLOR_HEX["black"]
            wb["color_name"] = "black"
            continue

        # --- Background subtraction via border pixels ---
        # Sample background from the 2px border ring of the crop
        if ch > 6 and cw > 6:
            border = 2
            bg_top = crop_hsv[:border, :].reshape(-1, 3)
            bg_bot = crop_hsv[-border:, :].reshape(-1, 3)
            bg_lft = crop_hsv[border:-border, :border].reshape(-1, 3)
            bg_rgt = crop_hsv[border:-border, -border:].reshape(-1, 3)
            bg_pixels = np.vstack([bg_top, bg_bot, bg_lft, bg_rgt])

            bg_med_h = float(np.median(bg_pixels[:, 0]))
            bg_med_s = float(np.median(bg_pixels[:, 1]))

            # If background is tinted (S > 15), remove text pixels
            # with similar hue to avoid false colored detections
            if bg_med_s > 15:
                hue_diff = np.minimum(
                    np.abs(text_pixels[:, 0].astype(float) - bg_med_h),
                    180.0 - np.abs(text_pixels[:, 0].astype(float) - bg_med_h),
                )
                keep = hue_diff > 20
                if np.any(keep):
                    text_pixels = text_pixels[keep]

        if len(text_pixels) < 3:
            wb["color"] = _COLOR_HEX["black"]
            wb["color_name"] = "black"
            continue

        # --- Classification using MEDIAN (robust to outliers) ---
        median_sat = float(np.median(text_pixels[:, 1]))
        sat_count = int(np.sum(text_pixels[:, 1] > sat_threshold))
        sat_ratio = sat_count / len(text_pixels)

        if median_sat < sat_threshold or sat_ratio < min_sat_ratio:
            wb["color"] = _COLOR_HEX["black"]
            wb["color_name"] = "black"
        else:
            # Use median hue of saturated pixels only for cleaner signal
            sat_pixels = text_pixels[text_pixels[:, 1] > sat_threshold]
            median_hue = float(np.median(sat_pixels[:, 0]))
            name = _hue_to_color_name(median_hue)

            # Red requires higher saturation — scanner artifacts on black
            # text often produce a slight warm tint (hue ~0) with low
            # saturation that would otherwise be misclassified as red.
            if name == "red" and median_sat < 90:
                wb["color"] = _COLOR_HEX["black"]
                wb["color_name"] = "black"
                continue

            wb["color"] = _COLOR_HEX.get(name, _COLOR_HEX["black"])
            wb["color_name"] = name
            colored_count += 1

    if colored_count:
        logger.info("color annotation: %d / %d words are colored",
                     colored_count, len(word_boxes))


# ---------------------------------------------------------------------------
# 2.  Recover colored text that OCR missed
# ---------------------------------------------------------------------------

def recover_colored_text(
    img_bgr: np.ndarray,
    existing_words: List[Dict],
    min_area: int = 40,
    max_regions: int = 60,
) -> List[Dict]:
    """Find colored text regions not covered by any existing word box.

    Returns a list of recovered word dicts with ``color``, ``color_name``,
    and ``recovered=True`` fields.  The ``text`` is set via a lightweight
    shape heuristic (e.g. ``!`` for tall narrow shapes) or ``?``.
    """
    if img_bgr is None:
        return []

    img_hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
    ih, iw = img_bgr.shape[:2]
    max_area = int(ih * iw * 0.005)

    # --- Build occupancy mask from existing words (adaptive padding) ---
    # Pad word boxes generously to prevent colored-pixel artifacts in
    # narrow inter-word gaps from being recovered as false characters.
    heights = [wb["height"] for wb in existing_words if wb.get("height", 0) > 0]
    median_h = int(np.median(heights)) if heights else 20
    pad = max(8, int(median_h * 0.35))

    occupied = np.zeros((ih, iw), dtype=np.uint8)
    for wb in existing_words:
        x1 = max(0, int(wb["left"]) - pad)
        y1 = max(0, int(wb["top"]) - pad)
        x2 = min(iw, int(wb["left"] + wb["width"]) + pad)
        y2 = min(ih, int(wb["top"] + wb["height"]) + pad)
        occupied[y1:y2, x1:x2] = 255

    recovered: List[Dict] = []

    for color_name, ranges in _COLOR_RANGES.items():
        # Create mask for this color
        mask = np.zeros((ih, iw), dtype=np.uint8)
        for lower, upper in ranges:
            mask = cv2.bitwise_or(mask, cv2.inRange(img_hsv, lower, upper))

        # Remove pixels already covered by existing OCR words
        mask = cv2.bitwise_and(mask, cv2.bitwise_not(occupied))

        # Morphological cleanup:
        # - Close with tall kernel to merge ! stroke + dot
        # - Open to remove noise specks
        kernel_close = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 8))
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close)
        kernel_open = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open)

        contours, _ = cv2.findContours(
            mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE,
        )

        candidates = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area < min_area or area > max_area:
                continue
            bx, by, bw, bh = cv2.boundingRect(cnt)
            if bh < 6:
                continue
            # Reject regions too wide to be single characters
            if bw > median_h * 4:
                continue
            candidates.append((area, bx, by, bw, bh))

        # Keep largest first, limited count
        candidates.sort(key=lambda c: c[0], reverse=True)

        for area, bx, by, bw, bh in candidates[:max_regions]:
            text = _identify_shape(bw, bh)
            recovered.append({
                "text": text,
                "left": bx,
                "top": by,
                "width": bw,
                "height": bh,
                "conf": 45,
                "color": _COLOR_HEX.get(color_name, "#000000"),
                "color_name": color_name,
                "recovered": True,
            })

    if recovered:
        logger.info(
            "color recovery: %d colored regions found (%s)",
            len(recovered),
            ", ".join(
                f"{c}: {sum(1 for r in recovered if r['color_name'] == c)}"
                for c in sorted({r["color_name"] for r in recovered})
            ),
        )

    return recovered


def _identify_shape(w: int, h: int) -> str:
    """Simple shape heuristic for common single-character text markers."""
    aspect = w / h if h > 0 else 1.0
    if aspect < 0.55 and h > 10:
        # Tall, narrow — likely exclamation mark
        return "!"
    if 0.6 < aspect < 1.5 and max(w, h) < 25:
        # Small, roughly square — bullet or dot
        return "•"
    return "?"
