"""Teachable Machine (Keras export) helper for ball localisation.

Export an Image Project from https://teachablemachine.withgoogle.com/
as Tensorflow / Keras, then drop into models/teachable/:

  keras_model.h5
  labels.txt

TM is a classifier, so we scan a coarse grid of crops and pick the
highest-confidence non-background cells as detections.
"""

from __future__ import annotations

from pathlib import Path
import threading

import cv2
import numpy as np

from .base import Detection

DEFAULT_MODEL_DIR = Path(__file__).resolve().parent.parent / "models" / "teachable"
INPUT_SIZE = 224


def _parse_labels(path: Path) -> list[str]:
    labels: list[str] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        # Teachable Machine lines look like: "0 Ball" or just "Ball"
        parts = line.split(maxsplit=1)
        if len(parts) == 2 and parts[0].isdigit():
            labels.append(parts[1].strip())
        else:
            labels.append(line)
    return labels


def _pick_target_index(labels: list[str]) -> int:
    """Prefer a class whose name looks like the object, else class 0."""
    for i, name in enumerate(labels):
        low = name.lower()
        if any(tok in low for tok in ("ball", "target", "object", "positive")):
            return i
    for i, name in enumerate(labels):
        low = name.lower()
        if "background" in low or "nothing" in low or "empty" in low or "neg" in low:
            continue
        return i
    return 0


class TeachableMachineModel:
    """Lazy-loaded Teachable Machine Keras classifier + grid search."""

    def __init__(
        self,
        model_dir: str | Path | None = None,
        *,
        confidence: float = 0.70,
        grid_cols: int = 5,
        grid_rows: int = 4,
        window_scale: float = 0.45,
    ) -> None:
        self._lock = threading.Lock()
        self.model_dir = Path(model_dir) if model_dir else DEFAULT_MODEL_DIR
        self.confidence = float(confidence)
        self.grid_cols = int(grid_cols)
        self.grid_rows = int(grid_rows)
        self.window_scale = float(window_scale)
        self._model = None
        self._labels: list[str] = []
        self._target_index = 0
        self._load_error: str | None = None

    @property
    def labels(self) -> list[str]:
        with self._lock:
            return list(self._labels)

    @property
    def load_error(self) -> str | None:
        with self._lock:
            return self._load_error

    @property
    def ready(self) -> bool:
        with self._lock:
            return self._model is not None

    def set_model_dir(self, model_dir: str | Path) -> None:
        with self._lock:
            self.model_dir = Path(model_dir)
            self._model = None
            self._labels = []
            self._load_error = None

    def ensure_loaded(self) -> bool:
        with self._lock:
            if self._model is not None:
                return True
            try:
                self._load_unlocked()
                return True
            except Exception as exc:  # noqa: BLE001 - surface any load failure in UI
                self._model = None
                self._load_error = str(exc)
                return False

    def _load_unlocked(self) -> None:
        try:
            import tensorflow as tf  # type: ignore
            from tensorflow import keras  # type: ignore
        except ImportError as exc:
            raise RuntimeError(
                "tensorflow not installed. Run: "
                "pip install tensorflow  (or tensorflow-cpu)"
            ) from exc

        model_path = self.model_dir / "keras_model.h5"
        labels_path = self.model_dir / "labels.txt"
        if not model_path.is_file():
            raise FileNotFoundError(
                f"missing {model_path.name} in {self.model_dir}. "
                "Export Keras from Teachable Machine and copy it here."
            )
        if not labels_path.is_file():
            raise FileNotFoundError(f"missing {labels_path.name} in {self.model_dir}")

        # Older TM exports pass a legacy `groups=` kwarg into DepthwiseConv2D.
        class _DepthwiseConv2D(keras.layers.DepthwiseConv2D):
            def __init__(self, *args, **kwargs):
                kwargs.pop("groups", None)
                super().__init__(*args, **kwargs)

        self._model = keras.models.load_model(
            model_path,
            compile=False,
            custom_objects={"DepthwiseConv2D": _DepthwiseConv2D},
        )
        self._labels = _parse_labels(labels_path)
        if not self._labels:
            raise ValueError(f"no labels found in {labels_path}")
        self._target_index = _pick_target_index(self._labels)
        self._load_error = None
        # Silence TF predict progress spam
        try:
            tf.get_logger().setLevel("ERROR")
        except Exception:
            pass

    @staticmethod
    def preprocess_bgr_crops(crops_bgr: list[np.ndarray]) -> np.ndarray:
        """Teachable Machine normalisation: RGB, 224x224, [-1, 1]."""
        batch = np.empty((len(crops_bgr), INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32)
        for i, crop in enumerate(crops_bgr):
            rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
            resized = cv2.resize(rgb, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_AREA)
            batch[i] = (resized.astype(np.float32) / 127.5) - 1.0
        return batch

    def detect(self, frame_bgr: np.ndarray, max_detections: int = 3) -> list[Detection]:
        if not self.ensure_loaded():
            return []

        with self._lock:
            model = self._model
            labels = list(self._labels)
            target_index = self._target_index
            conf_min = self.confidence
            cols = self.grid_cols
            rows = self.grid_rows
            window_scale = self.window_scale

        assert model is not None
        h, w = frame_bgr.shape[:2]
        win = int(round(min(h, w) * window_scale))
        win = max(48, min(win, h, w))

        crops: list[np.ndarray] = []
        centres: list[tuple[float, float, float]] = []  # cx, cy, radius
        for r in range(rows):
            for c in range(cols):
                cx = (c + 0.5) * w / cols
                cy = (r + 0.5) * h / rows
                x0 = int(round(cx - win / 2))
                y0 = int(round(cy - win / 2))
                x0 = max(0, min(x0, w - win))
                y0 = max(0, min(y0, h - win))
                crop = frame_bgr[y0 : y0 + win, x0 : x0 + win]
                if crop.shape[0] < 16 or crop.shape[1] < 16:
                    continue
                crops.append(crop)
                centres.append((float(x0 + win / 2), float(y0 + win / 2), float(win * 0.45)))

        if not crops:
            return []

        batch = self.preprocess_bgr_crops(crops)
        preds = model.predict(batch, verbose=0)
        preds = np.asarray(preds)

        found: list[Detection] = []
        target_name = labels[target_index] if target_index < len(labels) else "target"
        for i, (cx, cy, radius) in enumerate(centres):
            scores = preds[i]
            score = float(scores[target_index])
            if score < conf_min:
                continue
            found.append(
                Detection(
                    label=f"tm:{target_name}",
                    cx=cx,
                    cy=cy,
                    radius=radius,
                    score=score,
                )
            )

        found.sort(key=lambda d: d.score, reverse=True)
        # Keep spatially spread detections (simple NMS by distance)
        kept: list[Detection] = []
        min_dist = win * 0.55
        for d in found:
            if any(
                (d.cx - k.cx) ** 2 + (d.cy - k.cy) ** 2 < min_dist * min_dist
                for k in kept
            ):
                continue
            kept.append(d)
            if len(kept) >= max_detections:
                break
        return kept