"""BoT-SORT tracker wrapper using Ultralytics' built-in tracker backend."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Sequence

import numpy as np
import yaml
from ultralytics import YOLO

from src.tracking.tracker import BaseTracker, TrackedObject


class BotSortTracker(BaseTracker):
    """Runs pretrained YOLO detection and BoT-SORT ID association per frame."""

    def __init__(self, detector_config_path: str | Path, tracker_config_path: str | Path) -> None:
        self.detector_config_path = Path(detector_config_path)
        self.tracker_config_path = Path(tracker_config_path)
        self.detector_config = self._load_config(self.detector_config_path)
        self.model = YOLO(self.detector_config.get("model_name", "yolo11n.pt"))
        self.names = self.model.names
        self.class_ids = self._resolve_class_filter(self.detector_config.get("class_filter"))

    @staticmethod
    def _load_config(path: Path) -> dict[str, Any]:
        with path.open("r", encoding="utf-8") as file:
            return yaml.safe_load(file) or {}

    def _resolve_class_filter(self, class_filter: Sequence[str | int] | None) -> list[int] | None:
        if not class_filter:
            return None
        name_to_id = {str(name).lower(): int(idx) for idx, name in self.names.items()}
        resolved: list[int] = []
        for item in class_filter:
            if isinstance(item, int):
                resolved.append(item)
                continue
            class_id = name_to_id.get(str(item).lower())
            if class_id is not None:
                resolved.append(class_id)
        return resolved or None

    def update(self, frame: np.ndarray, frame_index: int) -> Sequence[TrackedObject]:
        results = self.model.track(
            source=frame,
            persist=True,
            tracker=str(self.tracker_config_path),
            conf=float(self.detector_config.get("confidence_threshold", 0.35)),
            iou=float(self.detector_config.get("iou_threshold", 0.5)),
            imgsz=int(self.detector_config.get("image_size", 1280)),
            device=self.detector_config.get("device"),
            classes=self.class_ids,
            agnostic_nms=bool(self.detector_config.get("agnostic_nms", False)),
            verbose=False,
        )

        tracked: list[TrackedObject] = []
        for result in results:
            if result.boxes is None or result.boxes.id is None:
                continue

            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy().astype(int)
            track_ids = result.boxes.id.cpu().numpy().astype(int)

            for bbox, confidence, class_id, track_id in zip(
                boxes, confidences, class_ids, track_ids, strict=False
            ):
                tracked.append(
                    TrackedObject(
                        id=int(track_id),
                        class_name=str(self.names.get(int(class_id), class_id)),
                        confidence=float(confidence),
                        bbox=tuple(float(v) for v in bbox),
                        frame_index=frame_index,
                    )
                )
        return tracked