import logging
from typing import Sequence

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

from models.detectors.base import DetectionResult, ObjectDetector


class HuggingFaceYoloV12BotSortDetector(ObjectDetector):
    """YOLOv12 model (BoT-SORT + ReID) hosted on Hugging Face."""

    REPO_ID = "wish44165/YOLOv12-BoT-SORT-ReID"
    WEIGHT_FILE = "MOT_yolov12n.pt"

    def __init__(self, score_threshold: float = 0.3) -> None:
        self.name = "hf_yolov12_bot_sort"
        self.score_threshold = score_threshold
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        logging.info(
            "Loading Hugging Face YOLOv12 BoT-SORT weights %s/%s onto %s",
            self.REPO_ID,
            self.WEIGHT_FILE,
            self.device,
        )
        weight_path = hf_hub_download(repo_id=self.REPO_ID, filename=self.WEIGHT_FILE)
        self.model = YOLO(weight_path)
        self.model.to(self.device)
        self.class_names = self.model.names

    def predict(self, frame: np.ndarray, queries: Sequence[str]) -> DetectionResult:
        device_arg = 0 if self.device.startswith("cuda") else "cpu"
        results = self.model.predict(
            source=frame,
            device=device_arg,
            conf=self.score_threshold,
            verbose=False,
        )
        result = results[0]
        boxes = result.boxes
        if boxes is None or boxes.xyxy is None:
            empty = np.empty((0, 4), dtype=np.float32)
            return DetectionResult(empty, [], [], [])

        xyxy = boxes.xyxy.cpu().numpy()
        scores = boxes.conf.cpu().numpy().tolist()
        label_ids = boxes.cls.cpu().numpy().astype(int).tolist()
        label_names = [self.class_names.get(idx, f"class_{idx}") for idx in label_ids]
        return DetectionResult(
            boxes=xyxy,
            scores=scores,
            labels=label_ids,
            label_names=label_names,
        )