"""
Wrap a trained YOLO model into structured, VLM-ready detections.

This is the DETERMINISTIC half of the pipeline. It:
  1. runs YOLO on an image
  2. returns clean structured detections (class, confidence, box)
  3. does simple, auditable PPE-compliance reasoning:
        "a head is detected but no helmet overlaps it" -> possible missing helmet

The VLM never decides compliance on its own — it is handed these facts. This keeps
the safety-critical judgement in code that can be logged and reviewed, in line with
the AiGiS design principle.
"""
from __future__ import annotations

from dataclasses import dataclass, asdict, field
from pathlib import Path

import yaml
from ultralytics import YOLO


@dataclass
class Detection:
    cls: str
    confidence: float
    box_xywhn: list[float]   # normalized [cx, cy, w, h] in [0,1]

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class ComplianceFinding:
    body_part: str
    expected_ppe: str
    status: str              # "present" | "missing" | "uncertain"
    detail: str

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class DetectionResult:
    image: str
    detections: list[Detection] = field(default_factory=list)
    compliance: list[ComplianceFinding] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "image": self.image,
            "detections": [d.to_dict() for d in self.detections],
            "compliance": [c.to_dict() for c in self.compliance],
        }


def _xywhn_to_corners(b: list[float]) -> tuple[float, float, float, float]:
    cx, cy, w, h = b
    return cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2


def _iou(a: list[float], b: list[float]) -> float:
    ax1, ay1, ax2, ay2 = _xywhn_to_corners(a)
    bx1, by1, bx2, by2 = _xywhn_to_corners(b)
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    if inter == 0:
        return 0.0
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    return inter / (area_a + area_b - inter)


class SafetyDetector:
    def __init__(self, weights: str, data_yaml: str, conf: float = 0.25):
        self.model = YOLO(weights)
        self.conf = conf
        with open(data_yaml) as f:
            cfg = yaml.safe_load(f)
        # body_part -> expected PPE class name
        self.ppe_pairs: dict[str, str] = cfg.get("ppe_pairs", {})

    def detect(self, image_path: str) -> DetectionResult:
        _, det_result = self._run(image_path)
        return det_result

    def detect_and_annotate(self, image, name: str | None = None):
        """
        Same as detect(), but also returns a BGR numpy array with the detection
        boxes drawn on it (via Ultralytics' own .plot()). Runs YOLO only once,
        so callers can get both the JSON and the annotated image from one call.

        `image` can be a file path OR an in-memory BGR array (e.g. from
        cv2.imdecode on an uploaded file) — Ultralytics accepts both.
        """
        result, det_result = self._run(image, name=name)
        annotated = result.plot()
        return det_result, annotated

    def _run(self, image, name: str | None = None):
        result = self.model(image, conf=self.conf, verbose=False)[0]
        names = result.names

        dets: list[Detection] = []
        for box in result.boxes:
            cls_name = names[int(box.cls)]
            dets.append(
                Detection(
                    cls=cls_name,
                    confidence=round(float(box.conf), 3),
                    box_xywhn=[round(float(x), 4) for x in box.xywhn[0].tolist()],
                )
            )

        compliance = self._assess_compliance(dets)
        if name is None:
            name = str(Path(image).name) if isinstance(image, (str, Path)) else "image"
        det_result = DetectionResult(image=name, detections=dets, compliance=compliance)
        return result, det_result

    def _assess_compliance(self, dets: list[Detection]) -> list[ComplianceFinding]:
        """
        For each body part that should be protected, check whether the expected
        PPE was detected overlapping it. Purely rule-based and explainable.
        """
        findings: list[ComplianceFinding] = []
        by_cls: dict[str, list[Detection]] = {}
        for d in dets:
            by_cls.setdefault(d.cls, []).append(d)

        for body_part, expected in self.ppe_pairs.items():
            parts = by_cls.get(body_part, [])
            if not parts:
                continue  # body part not visible -> nothing to assess
            ppe_items = by_cls.get(expected, [])
            for part in parts:
                overlap = max(
                    (_iou(part.box_xywhn, p.box_xywhn) for p in ppe_items),
                    default=0.0,
                )
                if overlap >= 0.10:
                    status, detail = "present", (
                        f"{expected} overlaps detected {body_part} (IoU {overlap:.2f})"
                    )
                elif ppe_items:
                    status, detail = "uncertain", (
                        f"{expected} present in image but not clearly over this "
                        f"{body_part} (max IoU {overlap:.2f})"
                    )
                else:
                    status, detail = "missing", (
                        f"{body_part} detected but no {expected} found in image"
                    )
                findings.append(
                    ComplianceFinding(body_part, expected, status, detail)
                )
        return findings