from pathlib import Path

from numpy import ndarray
from pydantic import BaseModel
from ultralytics import YOLO


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class Polygon(BaseModel):
    cls_id: int
    conf: float
    points: list[tuple[int, int]]


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list[Polygon] | None = None
    keypoints: list[tuple[int, int]] | None = None


# Element: manak0/Detect-road-signs — objects: ["road sign"] (single class, id 0).
# Trained on LVIS street_sign + signboard + stop_sign, all mapped to class 0.
# CONF_THRESHOLD is set by the sweep that replicates the element's own scorer
# (0.6*mAP50 + 0.4*(1 - FP_per_image/10)).
CONF_THRESHOLD = 0.40
IMGSZ = 1280  # matches element preproc resize_long


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        self.model = YOLO(str(path_hf_repo / "road-sign-detection.pt"))
        print("✅ Road-sign model loaded")

    def __repr__(self) -> str:
        return f"Detect-road-signs miner: {type(self.model).__name__} @ imgsz={IMGSZ}, conf={CONF_THRESHOLD}"

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results = self.model.predict(
            batch_images,
            imgsz=IMGSZ,
            conf=CONF_THRESHOLD,
            verbose=False,
            half=True,
        )
        out: list[TVFrameResult] = []
        for i, det in enumerate(results):
            boxes: list[BoundingBox] = []
            if det.boxes is not None:
                for box in det.boxes.data:
                    x1, y1, x2, y2, conf, cls_id = box.tolist()
                    boxes.append(
                        BoundingBox(
                            x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
                            cls_id=int(cls_id), conf=float(conf),
                        )
                    )
            out.append(TVFrameResult(frame_id=offset + i, boxes=boxes))
        return out