"""TurboVision public-track miner for manak0/Detect-worksite-ppe-compliance-detection-probe-set.

Model: YOLO26n exported to ONNX (input 1x3x704x704, NMS embedded, output 1x300x6
rows of x1, y1, x2, y2, score, class in letterboxed pixel space).

cls_id in the response is the index into the element's manifest `objects` list.
"""

from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
from numpy import ndarray
from pydantic import BaseModel

MANIFEST_CLASSES = [
    "safety helmet",
    "high-visibility safety vest",
    "protective gloves",
    "protective eyewear",
]

# Per-class confidence thresholds, indexed like MANIFEST_CLASSES.
CLASS_CONF = np.array([0.40, 0.40, 0.40, 0.30], dtype=np.float32)

MODEL_FILENAME = "weights.onnx"
LETTERBOX_FILL = 114


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


def _letterbox(image: ndarray, size: int) -> tuple[ndarray, float, int, int]:
    """Ultralytics-style centred letterbox to a size x size canvas."""
    height, width = image.shape[:2]
    ratio = min(size / height, size / width)
    new_w = int(round(width * ratio))
    new_h = int(round(height * ratio))
    if (new_w, new_h) != (width, height):
        image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    pad_w = (size - new_w) / 2
    pad_h = (size - new_h) / 2
    top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
    left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
    canvas = cv2.copyMakeBorder(
        image,
        top,
        bottom,
        left,
        right,
        cv2.BORDER_CONSTANT,
        value=(LETTERBOX_FILL, LETTERBOX_FILL, LETTERBOX_FILL),
    )
    return canvas, ratio, left, top


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        model_path = Path(path_hf_repo) / MODEL_FILENAME
        if not model_path.is_file():
            raise FileNotFoundError(f"{MODEL_FILENAME} not found in {path_hf_repo}")

        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        options.intra_op_num_threads = 2
        options.inter_op_num_threads = 1

        providers = ["CPUExecutionProvider"]
        if "CUDAExecutionProvider" in ort.get_available_providers():
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        try:
            self.session = ort.InferenceSession(
                str(model_path), sess_options=options, providers=providers
            )
        except Exception as exc:
            print(f"CUDA session failed, falling back to CPU: {exc}")
            self.session = ort.InferenceSession(
                str(model_path), sess_options=options, providers=["CPUExecutionProvider"]
            )

        model_input = self.session.get_inputs()[0]
        self.input_name = model_input.name
        self.input_size = int(model_input.shape[2])
        self.output_name = self.session.get_outputs()[0].name
        self._check_class_order()

        warmup = np.zeros((self.input_size, self.input_size, 3), dtype=np.uint8)
        for _ in range(2):
            self._infer_one(warmup)
        print(f"Miner ready: providers={self.session.get_providers()} input={self.input_size}")

    def __repr__(self) -> str:
        return (
            f"PPE YOLO26n ONNX input={self.input_size} "
            f"providers={self.session.get_providers()}"
        )

    def _check_class_order(self) -> None:
        names = self.session.get_modelmeta().custom_metadata_map.get("names", "")
        for idx, name in enumerate(MANIFEST_CLASSES):
            if f"{idx}: '{name}'" not in names:
                raise ValueError(f"ONNX class order does not match manifest: {names}")

    def _infer_one(self, image: ndarray) -> list[BoundingBox]:
        canvas, ratio, pad_x, pad_y = _letterbox(image, self.input_size)
        rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
        tensor = np.ascontiguousarray(rgb.transpose(2, 0, 1)[None], dtype=np.float32) / 255.0
        raw = self.session.run([self.output_name], {self.input_name: tensor})[0][0]

        scores = raw[:, 4]
        cls_ids = raw[:, 5].astype(np.int64)
        known = (cls_ids >= 0) & (cls_ids < len(MANIFEST_CLASSES))
        keep = known & (scores >= CLASS_CONF[np.clip(cls_ids, 0, len(CLASS_CONF) - 1)])
        if not keep.any():
            return []

        boxes = raw[keep, :4].astype(np.float64)
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_x) / ratio
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_y) / ratio
        height, width = image.shape[:2]
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, width)
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, height)

        results: list[BoundingBox] = []
        for (x1, y1, x2, y2), score, cls_id in zip(boxes, scores[keep], cls_ids[keep]):
            x1, y1, x2, y2 = (int(round(v)) for v in (x1, y1, x2, y2))
            if x2 <= x1 or y2 <= y1:
                continue
            results.append(
                BoundingBox(
                    x1=x1, y1=y1, x2=x2, y2=y2, cls_id=int(cls_id), conf=float(score)
                )
            )
        return results

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        for index, image in enumerate(batch_images):
            try:
                boxes = self._infer_one(image)
            except Exception as exc:
                print(f"Inference failed on frame {offset + index}: {exc}")
                boxes = []
            results.append(
                TVFrameResult(
                    frame_id=offset + index,
                    boxes=boxes,
                    polygons=[],
                    keypoints=[(0, 0)] * max(0, int(n_keypoints)),
                )
            )
        return results