"""HSV colour mask + swappable ball finders (contour / hough / blob / TM)."""

from __future__ import annotations

from pathlib import Path
import threading

import cv2
import numpy as np

from .base import Detection, TargetDetector
from .teachable_machine import TeachableMachineModel

METHODS = ("contour", "hough", "blob", "rect", "teachable")
METHOD_LABELS = {
    "contour": "Contour",
    "hough": "Hough circles",
    "blob": "Blob",
    "rect": "Rectangle",
    "teachable": "Teachable Machine",
}


class OrangeBallDetector(TargetDetector):
    name = "orange_ball"

    def __init__(
        self,
        *,
        # Centre colour in OpenCV HSV (H: 0-179)
        hsv_center: tuple[int, int, int] = (15, 180, 200),
        h_tol: int = 10,
        s_tol: int = 80,
        v_tol: int = 80,
        min_area: float = 80.0,
        min_circularity: float = 0.55,
        min_rectangularity: float = 0.65,
        max_detections: int = 3,
        method: str = "contour",
        teachable_dir: str | Path | None = None,
    ) -> None:
        self._lock = threading.Lock()
        self.hsv_center = np.array(hsv_center, dtype=np.int32)
        self.h_tol = int(h_tol)
        self.s_tol = int(s_tol)
        self.v_tol = int(v_tol)
        self.min_area = float(min_area)
        self.min_circularity = float(min_circularity)
        self.min_rectangularity = float(min_rectangularity)
        self.max_detections = max_detections
        self.method = self._normalize_method(method)
        self._blob_detector = self._make_blob_detector(
            self.min_area, self.min_circularity
        )
        self.teachable = TeachableMachineModel(teachable_dir)

    @staticmethod
    def _normalize_method(method: str) -> str:
        key = method.strip().lower()
        if key not in METHODS:
            raise ValueError(f"unknown method '{method}'. known: {', '.join(METHODS)}")
        return key

    @staticmethod
    def _make_blob_detector(min_area: float, min_circularity: float):
        params = cv2.SimpleBlobDetector_Params()
        params.filterByColor = True
        params.blobColor = 255
        params.filterByArea = True
        params.minArea = float(max(1.0, min_area))
        params.maxArea = 1e6
        params.filterByCircularity = True
        params.minCircularity = float(max(0.01, min(1.0, min_circularity)))
        params.filterByInertia = True
        params.minInertiaRatio = 0.4
        params.filterByConvexity = True
        params.minConvexity = 0.6
        params.minThreshold = 50
        params.maxThreshold = 255
        params.thresholdStep = 50
        return cv2.SimpleBlobDetector_create(params)

    def get_hsv_center(self) -> tuple[int, int, int]:
        with self._lock:
            h, s, v = (int(x) for x in self.hsv_center)
            return h, s, v

    def get_bgr_center(self) -> tuple[int, int, int]:
        h, s, v = self.get_hsv_center()
        pixel = np.uint8([[[h, s, v]]])
        bgr = cv2.cvtColor(pixel, cv2.COLOR_HSV2BGR)[0, 0]
        return int(bgr[0]), int(bgr[1]), int(bgr[2])

    def get_tolerance(self) -> tuple[int, int, int]:
        with self._lock:
            return self.h_tol, self.s_tol, self.v_tol

    def get_method(self) -> str:
        with self._lock:
            return self.method

    def set_method(self, method: str) -> None:
        key = self._normalize_method(method)
        with self._lock:
            self.method = key
        if key == "teachable":
            self.teachable.ensure_loaded()

    def set_teachable_dir(self, model_dir: str | Path) -> None:
        self.teachable.set_model_dir(model_dir)
        if self.get_method() == "teachable":
            self.teachable.ensure_loaded()

    def teachable_status(self) -> str:
        if self.teachable.ready:
            labels = ", ".join(self.teachable.labels) or "?"
            return f"ready ({labels})"
        err = self.teachable.load_error
        if err:
            return f"error: {err}"
        return f"not loaded - put keras_model.h5 + labels.txt in {self.teachable.model_dir}"

    def set_from_bgr(self, b: int, g: int, r: int) -> None:
        pixel = np.uint8([[[b, g, r]]])
        hsv = cv2.cvtColor(pixel, cv2.COLOR_BGR2HSV)[0, 0]
        self.set_hsv_center(int(hsv[0]), int(hsv[1]), int(hsv[2]))

    def set_hsv_center(self, h: int, s: int, v: int) -> None:
        with self._lock:
            self.hsv_center = np.array(
                [np.clip(h, 0, 179), np.clip(s, 0, 255), np.clip(v, 0, 255)],
                dtype=np.int32,
            )

    def set_tolerance(
        self,
        h_tol: int | None = None,
        s_tol: int | None = None,
        v_tol: int | None = None,
    ) -> None:
        with self._lock:
            if h_tol is not None:
                self.h_tol = int(np.clip(h_tol, 1, 90))
            if s_tol is not None:
                self.s_tol = int(np.clip(s_tol, 1, 255))
            if v_tol is not None:
                self.v_tol = int(np.clip(v_tol, 1, 255))
            self._blob_detector = self._make_blob_detector(
                self.min_area, self.min_circularity
            )

    def _snapshot(
        self,
    ) -> tuple[np.ndarray, np.ndarray, float, float, float, int, str, int, int]:
        with self._lock:
            h, s, v = (int(x) for x in self.hsv_center)
            lo = np.array(
                [
                    max(0, h - self.h_tol),
                    max(0, s - self.s_tol),
                    max(0, v - self.v_tol),
                ],
                dtype=np.uint8,
            )
            hi = np.array(
                [
                    min(179, h + self.h_tol),
                    min(255, s + self.s_tol),
                    min(255, v + self.v_tol),
                ],
                dtype=np.uint8,
            )
            return (
                lo,
                hi,
                self.min_area,
                self.min_circularity,
                self.min_rectangularity,
                self.max_detections,
                self.method,
                h,
                self.h_tol,
            )

    def build_mask(self, frame_bgr: np.ndarray) -> np.ndarray:
        hsv_low, hsv_high, *_rest, h_center, h_tol = self._snapshot()
        hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
        s_lo, s_hi = int(hsv_low[1]), int(hsv_high[1])
        v_lo, v_hi = int(hsv_low[2]), int(hsv_high[2])

        # Red hues wrap at 0/179 — merge two ranges when needed
        if h_center - h_tol < 0:
            mask_a = cv2.inRange(
                hsv,
                np.array([0, s_lo, v_lo], dtype=np.uint8),
                np.array([h_center + h_tol, s_hi, v_hi], dtype=np.uint8),
            )
            wrap_lo = 179 + (h_center - h_tol)
            mask_b = cv2.inRange(
                hsv,
                np.array([wrap_lo, s_lo, v_lo], dtype=np.uint8),
                np.array([179, s_hi, v_hi], dtype=np.uint8),
            )
            mask = cv2.bitwise_or(mask_a, mask_b)
        elif h_center + h_tol > 179:
            mask_a = cv2.inRange(
                hsv,
                np.array([h_center - h_tol, s_lo, v_lo], dtype=np.uint8),
                np.array([179, s_hi, v_hi], dtype=np.uint8),
            )
            wrap_hi = h_center + h_tol - 179
            mask_b = cv2.inRange(
                hsv,
                np.array([0, s_lo, v_lo], dtype=np.uint8),
                np.array([wrap_hi, s_hi, v_hi], dtype=np.uint8),
            )
            mask = cv2.bitwise_or(mask_a, mask_b)
        else:
            mask = cv2.inRange(hsv, hsv_low, hsv_high)

        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
        return mask

    def detect(self, frame_bgr: np.ndarray) -> list[Detection]:
        (
            _,
            _,
            min_area,
            min_circularity,
            min_rectangularity,
            max_detections,
            method,
            *_rest,
        ) = self._snapshot()
        if method == "teachable":
            found = self.teachable.detect(frame_bgr, max_detections=max_detections)
        else:
            mask = self.build_mask(frame_bgr)
            if method == "hough":
                found = self._detect_hough(mask, min_area, max_detections)
            elif method == "blob":
                found = self._detect_blob(mask, max_detections)
            elif method == "rect":
                found = self._detect_rect(
                    mask, min_area, min_rectangularity, max_detections
                )
            else:
                found = self._detect_contour(
                    mask, min_area, min_circularity, max_detections
                )
        found.sort(key=lambda d: d.score, reverse=True)
        return found[:max_detections]

    def _detect_contour(
        self,
        mask: np.ndarray,
        min_area: float,
        min_circularity: float,
        max_detections: int,
    ) -> list[Detection]:
        contours, _ = cv2.findContours(
            mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        found: list[Detection] = []
        for contour in contours:
            area = float(cv2.contourArea(contour))
            if area < min_area:
                continue
            peri = cv2.arcLength(contour, True)
            if peri <= 0:
                continue
            circularity = float(4.0 * np.pi * area / (peri * peri))
            if circularity < min_circularity:
                continue

            (cx, cy), radius = cv2.minEnclosingCircle(contour)
            score = max(
                0.0,
                min(1.0, 0.5 * circularity + 0.5 * min(1.0, area / 4000.0)),
            )
            found.append(
                Detection(
                    label="contour",
                    cx=float(cx),
                    cy=float(cy),
                    radius=float(radius),
                    score=score,
                )
            )
        return found[:max_detections]

    def _detect_rect(
        self,
        mask: np.ndarray,
        min_area: float,
        min_rectangularity: float,
        max_detections: int,
    ) -> list[Detection]:
        contours, _ = cv2.findContours(
            mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        found: list[Detection] = []
        for contour in contours:
            area = float(cv2.contourArea(contour))
            if area < min_area:
                continue
            peri = cv2.arcLength(contour, True)
            if peri <= 0:
                continue
            approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
            if len(approx) != 4:
                continue
            if not cv2.isContourConvex(approx):
                continue

            rect = cv2.minAreaRect(contour)
            (cx, cy), (w, h), angle = rect
            box_area = float(w) * float(h)
            if box_area <= 0:
                continue
            rectangularity = area / box_area
            if rectangularity < min_rectangularity:
                continue

            score = max(
                0.0,
                min(
                    1.0,
                    0.6 * rectangularity + 0.4 * min(1.0, area / 4000.0),
                ),
            )
            radius = max(float(w), float(h)) * 0.5
            found.append(
                Detection(
                    label="rect",
                    cx=float(cx),
                    cy=float(cy),
                    radius=radius,
                    score=score,
                    width=float(w),
                    height=float(h),
                    angle=float(angle),
                )
            )
        return found[:max_detections]

    def _detect_hough(
        self,
        mask: np.ndarray,
        min_area: float,
        max_detections: int,
    ) -> list[Detection]:
        # Estimate radius bounds from min_area of a filled circle
        min_radius = max(3, int(round(np.sqrt(min_area / np.pi))))
        blurred = cv2.GaussianBlur(mask, (9, 9), 2)
        circles = cv2.HoughCircles(
            blurred,
            cv2.HOUGH_GRADIENT,
            dp=1.2,
            minDist=max(12, min_radius * 2),
            param1=80,
            param2=18,
            minRadius=min_radius,
            maxRadius=0,
        )
        found: list[Detection] = []
        if circles is None:
            return found
        for x, y, r in circles[0]:
            area = float(np.pi * r * r)
            score = max(0.0, min(1.0, min(1.0, area / 4000.0)))
            found.append(
                Detection(
                    label="hough",
                    cx=float(x),
                    cy=float(y),
                    radius=float(r),
                    score=score,
                )
            )
        return found[:max_detections]

    def _detect_blob(self, mask: np.ndarray, max_detections: int) -> list[Detection]:
        with self._lock:
            detector = self._blob_detector
        keypoints = detector.detect(mask)
        found: list[Detection] = []
        for kp in keypoints:
            radius = float(kp.size) * 0.5
            area = float(np.pi * radius * radius)
            # OpenCV blob response is not always normalized; fold size in
            response = float(kp.response) if kp.response else 0.0
            score = max(
                0.0,
                min(1.0, 0.4 * min(1.0, response) + 0.6 * min(1.0, area / 4000.0)),
            )
            found.append(
                Detection(
                    label="blob",
                    cx=float(kp.pt[0]),
                    cy=float(kp.pt[1]),
                    radius=radius,
                    score=score,
                )
            )
        return found[:max_detections]