from __future__ import annotations

from enum import Enum
from typing import Optional

import numpy as np

from eye.detection.core import Detections


class TrackerType(str, Enum):
    SORT = "sort"
    BYTETRACK = "bytetrack"
    BOTSORT = "botsort"


def _get_tracker_backend(tracker_type: TrackerType, **kwargs):
    """Lazy-import and instantiate the requested tracker backend.

    This avoids top-level imports of filterpy / lap which may not be installed.
    """
    if tracker_type == TrackerType.SORT:
        from eye.core.trackers.sort_tracker import SORTTracker
        return SORTTracker(
            max_age=kwargs.get("max_age", 30),
            min_hits=kwargs.get("min_hits", 3),
            iou_threshold=kwargs.get("iou_threshold", 0.3),
        )
    elif tracker_type == TrackerType.BOTSORT:
        from eye.core.trackers.botsort_tracker import BoTSORTTracker
        return BoTSORTTracker(
            max_age=kwargs.get("max_age", 30),
            min_hits=kwargs.get("min_hits", 3),
            iou_threshold=kwargs.get("iou_threshold", 0.3),
            proximity_thresh=kwargs.get("proximity_thresh", 0.5),
            appearance_thresh=kwargs.get("appearance_thresh", 0.25),
        )
    else:
        from eye.core.trackers.bytetrack_tracker import ByteTrackTracker
        return ByteTrackTracker(
            max_age=kwargs.get("max_age", 30),
            min_hits=kwargs.get("min_hits", 3),
            iou_threshold=kwargs.get("iou_threshold", 0.3),
            high_threshold=kwargs.get("high_threshold", 0.5),
            low_threshold=kwargs.get("low_threshold", 0.1),
        )


def _inflate_xyxy(xyxy: np.ndarray, factor: float) -> np.ndarray:
    if factor is None or factor == 1.0:
        return xyxy
    xyxy = np.asarray(xyxy, dtype=np.float32)
    w = xyxy[:, 2] - xyxy[:, 0]
    h = xyxy[:, 3] - xyxy[:, 1]
    cx = xyxy[:, 0] + 0.5 * w
    cy = xyxy[:, 1] + 0.5 * h
    w2 = 0.5 * w * factor
    h2 = 0.5 * h * factor
    out = np.empty_like(xyxy, dtype=np.float32)
    out[:, 0] = cx - w2
    out[:, 1] = cy - h2
    out[:, 2] = cx + w2
    out[:, 3] = cy + h2
    return out


def _iou_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """IoU matrix between a (N,4) and b (M,4) in xyxy."""
    if a.size == 0 or b.size == 0:
        return np.zeros((a.shape[0], b.shape[0]), dtype=np.float32)

    a = a.astype(np.float32)
    b = b.astype(np.float32)

    xx1 = np.maximum(a[:, None, 0], b[None, :, 0])
    yy1 = np.maximum(a[:, None, 1], b[None, :, 1])
    xx2 = np.minimum(a[:, None, 2], b[None, :, 2])
    yy2 = np.minimum(a[:, None, 3], b[None, :, 3])

    inter_w = np.maximum(0.0, xx2 - xx1)
    inter_h = np.maximum(0.0, yy2 - yy1)
    inter = inter_w * inter_h

    area_a = (a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1])
    area_b = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])

    union = area_a[:, None] + area_b[None, :] - inter
    return inter / (union + 1e-6)


def _assign_ids_by_iou(
    det_xyxy: np.ndarray,
    trk_xyxy: np.ndarray,
    trk_ids: np.ndarray,
    min_iou: float,
) -> np.ndarray:
    """Greedy IoU assignment. Returns per-detection tracker_id, -1 if unmatched."""
    n = det_xyxy.shape[0]
    if n == 0 or trk_xyxy.shape[0] == 0:
        return np.full((n,), -1, dtype=np.int64)

    ious = _iou_matrix(det_xyxy, trk_xyxy)
    assigned = np.full((n,), -1, dtype=np.int64)

    used_tracks: set[int] = set()
    for det_i in range(n):
        best_j = int(np.argmax(ious[det_i]))
        best_iou = float(ious[det_i, best_j])
        if best_iou < min_iou:
            continue
        if best_j in used_tracks:
            continue
        used_tracks.add(best_j)
        assigned[det_i] = int(trk_ids[best_j])

    return assigned


class Tracker:
    """Universal tracker wrapper.

    Supports SORT, ByteTrack, and BoT-SORT backends with optional
    built-in smoothing and box inflation for robust tracking.

    For production use, you can also use ``eye.ByteTrack`` directly.
    """

    def __init__(
        self,
        tracker_type: TrackerType = TrackerType.BYTETRACK,
        inflation_factor: float = 1.0,
        max_age: int = 30,
        min_hits: int = 3,
        iou_threshold: float = 0.3,
        high_threshold: float = 0.5,
        low_threshold: float = 0.1,
        min_iou_assignment: float = 0.3,
        # BoT-SORT specific
        proximity_thresh: float = 0.5,
        appearance_thresh: float = 0.25,
        # Smoothing
        enable_smoothing: bool = False,
        smoothing_alpha: float = 0.3,
    ) -> None:
        self.tracker_type = TrackerType(tracker_type)
        self.inflation_factor = float(inflation_factor)
        self.min_iou_assignment = float(min_iou_assignment)
        self.enable_smoothing = enable_smoothing

        self._tracker = _get_tracker_backend(
            self.tracker_type,
            max_age=max_age,
            min_hits=min_hits,
            iou_threshold=iou_threshold,
            high_threshold=high_threshold,
            low_threshold=low_threshold,
            proximity_thresh=proximity_thresh,
            appearance_thresh=appearance_thresh,
        )

        # Smoothing setup (lazy import to avoid requiring filterpy just for EMA)
        self._smoother = None
        if enable_smoothing:
            try:
                from eye.detection.tools.smoothing import ExponentialSmoothing
                self._smoother = ExponentialSmoothing(alpha=smoothing_alpha)
            except ImportError:
                pass

    def update(self, detections: Detections) -> Detections:
        """Track detections and assign tracker IDs.

        Args:
            detections: Input :class:`Detections` with ``xyxy`` and ``confidence``.

        Returns:
            Filtered :class:`Detections` containing only matched tracks
            with ``tracker_id`` set.
        """
        if detections is None or len(detections) == 0:
            if detections is not None:
                detections.tracker_id = np.empty((0,), dtype=np.int64)
            return detections

        xyxy = np.asarray(detections.xyxy, dtype=np.float32)
        conf = getattr(detections, "confidence", None)
        if conf is None:
            conf = np.ones((xyxy.shape[0],), dtype=np.float32)
        conf = np.asarray(conf, dtype=np.float32).reshape(-1)

        inflated = _inflate_xyxy(xyxy, self.inflation_factor)
        dets_for_tracker = np.column_stack([inflated, conf])

        tracked = self._tracker.update(dets_for_tracker)
        if tracked is None or tracked.size == 0:
            detections.tracker_id = np.full((len(detections),), -1, dtype=np.int64)
            return detections[detections.tracker_id != -1]

        trk_xyxy = np.asarray(tracked[:, :4], dtype=np.float32)
        trk_ids = np.asarray(tracked[:, 4], dtype=np.int64)

        assigned = _assign_ids_by_iou(inflated, trk_xyxy, trk_ids, min_iou=self.min_iou_assignment)
        detections.tracker_id = assigned
        result = detections[detections.tracker_id != -1]

        # Apply smoothing if enabled
        if self._smoother is not None and len(result) > 0:
            try:
                from eye.detection.tools.smoothing import smooth_detections
                result = smooth_detections(result, self._smoother)
            except Exception:
                pass

        return result

    def reset(self) -> None:
        """Reset the tracker state, clearing all tracks."""
        if hasattr(self._tracker, "reset"):
            self._tracker.reset()
        else:
            # Fallback: clear known tracker internals
            if hasattr(self._tracker, "trackers"):
                self._tracker.trackers = []
            if hasattr(self._tracker, "frame_count"):
                self._tracker.frame_count = 0

        if self._smoother is not None and hasattr(self._smoother, "reset"):
            self._smoother.reset()
