"""Registration metadata for the built-in PPE YOLOv8 checkpoints."""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from ..base import DetectionModel
from ..registry import DetectionModelRegistry, DetectionModelSpec
from .yolo import UltralyticsYoloModel


CONSTRUCTION_PPE_LABELS = (
    "Boots",
    "Ear-Protection",
    "Glass",
    "Glove",
    "Hard_hat",
    "Mask",
    "No-Boots",
    "No-Ear-Protection",
    "No-Glass",
    "No-Glove",
    "No-Helmet",
    "No-Mask",
    "No-Vest",
    "Worker",
    "Vest",
    "Circular_Saw",
    "Fire_Extinguisher",
    "Fire_prevention_Net",
    "Welding_Equipment",
)

PPE_6CLASS_LABELS = (
    "Gloves",
    "Vest",
    "goggles",
    "helmet",
    "mask",
    "safety_shoe",
)


def _create_construction_ppe_model(
    params: Mapping[str, Any],
) -> DetectionModel:
    return UltralyticsYoloModel(
        params,
        expected_labels=CONSTRUCTION_PPE_LABELS,
    )


def _create_ppe_6classes_model(
    params: Mapping[str, Any],
) -> DetectionModel:
    return UltralyticsYoloModel(
        params,
        expected_labels=PPE_6CLASS_LABELS,
    )


def register_builtin_detection_models(registry: DetectionModelRegistry) -> None:
    registry.register(
        DetectionModelSpec(
            model_id="construction-ppe-yolov8@1",
            name="Construction PPE YOLOv8s",
            supported_labels=CONSTRUCTION_PPE_LABELS,
            factory=_create_construction_ppe_model,
            backend="ultralytics-yolo",
            description=(
                "Construction-site worker and personal protective equipment "
                "detector with positive and No-* violation labels"
            ),
        )
    )
    registry.register(
        DetectionModelSpec(
            model_id="ppe-6classes-yolov8n@1",
            name="PPE Detection YOLOv8n (6 Classes)",
            supported_labels=PPE_6CLASS_LABELS,
            factory=_create_ppe_6classes_model,
            backend="ultralytics-yolo",
            description=(
                "Lightweight six-class personal protective equipment detector "
                "for gloves, vests, goggles, helmets, masks and safety shoes"
            ),
        )
    )


__all__ = [
    "CONSTRUCTION_PPE_LABELS",
    "PPE_6CLASS_LABELS",
    "register_builtin_detection_models",
]
