"""
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, PER PERSON:
        "this person's head is detected and no helmet is on it" -> helmet missing

The VLM never decides compliance on its own. It is handed these facts, and the app
renders them as a table straight from this module. This keeps the safety-critical
judgement in code that can be logged and reviewed, in line with the AiGiS design
principle.

What is checked comes from the data YAML (sh17.yaml), not from this file:

    ppe_rules:
      - item: eye_protection          # name shown in the PPE table
        body_part: face               # detector class the item protects
        any_of: [glasses, face_guard] # ANY one of these satisfies the rule
        level: advisory               # "required" or "advisory"
    compliance_conf: 0.40             # confidence floor for compliance evidence

The older format is still accepted (every listed item is treated as required):

    ppe_pairs:
      face: [glasses, face_mask_medical]

How one person's item gets its status:
  present    a confident PPE box sits on the body part (or is on the person, for
             items such as gloves where the gloved hand is not detected as "hands")
  missing    the body part is confidently detected and no PPE box is on it.
             PPE worn by OTHER people never changes this.
  uncertain  the only evidence is a low-confidence box, or the PPE is on the
             person but not clearly over the body part
A body part that is not detected at all produces no finding (not assessed).
"""
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        # class name(s) that satisfy the rule, e.g. "glasses or face_guard"
    status: str              # "present" | "missing" | "uncertain"
    detail: str
    # --- added fields (older consumers can ignore them) ---
    item: str = ""           # rule name, e.g. "eye_protection"
    level: str = "required"  # "required" | "advisory"
    required: bool = True
    person_id: int | None = None   # 1..N, people numbered left to right; None = not tied to a person
    confidence: float | None = None  # confidence of the box that decided the status

    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)
    people: int = 0          # number of people the compliance check assessed

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


@dataclass
class PpeRule:
    item: str
    body_part: str
    any_of: list[str]
    level: str = "required"


# ---------------------------------------------------------------------------
# Geometry
# ---------------------------------------------------------------------------
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)


def _center_in(inner: list[float], outer: list[float]) -> bool:
    """True if the centre of `inner` lies inside the box `outer`."""
    x1, y1, x2, y2 = _xywhn_to_corners(outer)
    return x1 <= inner[0] <= x2 and y1 <= inner[1] <= y2


# ---------------------------------------------------------------------------
# Compliance settings
# ---------------------------------------------------------------------------
# Minimum IoU for a PPE box to count as "covering" a body-part box.
MATCH_IOU = 0.10

# Used when the YAML does not set compliance_conf.
DEFAULT_COMPLIANCE_CONF = 0.40

PERSON = "person"
LEVELS = ("required", "advisory")

# Body parts a person has two of. One protected hand says nothing about the other,
# so PPE elsewhere on the same person cannot soften a bare one to "uncertain".
PAIRED_PARTS = {"hands", "foot", "ear"}

_RANK = {"present": 0, "uncertain": 1, "missing": 2}


def load_ppe_rules(cfg: dict, class_names) -> list[PpeRule]:
    """Read ppe_rules (new format) or ppe_pairs (old format) from the YAML dict."""
    rules: list[PpeRule] = []
    raw_rules = cfg.get("ppe_rules")
    if raw_rules:
        for r in raw_rules:
            any_of = r["any_of"]
            any_of = [any_of] if isinstance(any_of, str) else list(any_of)
            level = str(r.get("level", "required")).lower()
            if level not in LEVELS:
                raise ValueError(
                    f"ppe_rules: level must be one of {LEVELS}, got '{level}'"
                )
            rules.append(
                PpeRule(
                    item=str(r.get("item", any_of[0])),
                    body_part=r["body_part"],
                    any_of=any_of,
                    level=level,
                )
            )
    else:
        # Old format: every listed item is its own required rule.
        for part, exp in (cfg.get("ppe_pairs") or {}).items():
            for e in ([exp] if isinstance(exp, str) else list(exp)):
                rules.append(PpeRule(item=e, body_part=part, any_of=[e]))

    known = set(class_names)
    for r in rules:
        unknown = [c for c in (r.body_part, *r.any_of) if c not in known]
        if unknown:
            raise ValueError(
                f"PPE rule '{r.item}' uses class name(s) not in the YAML names: {unknown}"
            )
    return rules


class SafetyDetector:
    def __init__(self, weights: str, data_yaml: str, conf: float = 0.25,
                 augment: bool = False, compliance_conf: float | None = None):
        self.model = YOLO(weights)
        self.conf = conf
        self.augment = augment   # test-time augmentation (slower, higher recall)
        with open(data_yaml) as f:
            cfg = yaml.safe_load(f)

        # Class names come from the YAML, NOT from the names embedded in the .pt.
        # The weights were trained with a mis-ordered names list, so the embedded
        # names are wrong. The label indices are correct; only the names were off.
        # The YAML (official SH17 index order) is the single source of truth.
        raw_names = cfg["names"]
        if isinstance(raw_names, dict):
            self.names = {int(k): v for k, v in raw_names.items()}
        else:
            self.names = dict(enumerate(raw_names))
        n_model = len(self.model.names)
        if len(self.names) != n_model:
            raise ValueError(
                f"{data_yaml} lists {len(self.names)} classes but the model has {n_model}"
            )
        self.model.model.names = self.names  # so annotated images use correct labels

        # What to check, and how sure a box must be to count as evidence.
        self.rules: list[PpeRule] = load_ppe_rules(cfg, self.names.values())
        if compliance_conf is None:
            compliance_conf = cfg.get("compliance_conf", DEFAULT_COMPLIANCE_CONF)
        self.compliance_conf = float(compliance_conf)

        # Kept for older code that reads detector.ppe_pairs.
        self.ppe_pairs: dict[str, list[str]] = {}
        for r in self.rules:
            self.ppe_pairs.setdefault(r.body_part, []).extend(r.any_of)

    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. Runs YOLO only once.
        `image` can be a file path OR an in-memory BGR array.
        """
        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, augment=self.augment,
                            verbose=False)[0]
        result.names = self.names  # guarantee correct labels in plot + JSON
        names = self.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)
        people = len(self._people(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, people=people)
        return result, det_result

    # -----------------------------------------------------------------------
    # Compliance
    # -----------------------------------------------------------------------
    def _people(self, dets: list[Detection]) -> list[Detection]:
        """Confidently detected people, numbered left to right (index 0 = person 1)."""
        people = [d for d in dets
                  if d.cls == PERSON and d.confidence >= self.compliance_conf]
        return sorted(people, key=lambda d: d.box_xywhn[0])

    def _assess_compliance(self, dets: list[Detection]) -> list[ComplianceFinding]:
        """
        For each person, and each PPE rule, decide present / missing / uncertain
        from the boxes that belong to THAT person. Purely rule-based and explainable.
        """
        people = self._people(dets)

        # Give every non-person box to the person whose box contains its centre
        # (the smallest such box if people overlap). None = belongs to nobody.
        groups: dict[int | None, list[Detection]] = {}
        for d in dets:
            if d.cls == PERSON:
                continue
            inside = [i for i, p in enumerate(people)
                      if _center_in(d.box_xywhn, p.box_xywhn)]
            if inside:
                i = min(inside,
                        key=lambda k: people[k].box_xywhn[2] * people[k].box_xywhn[3])
                pid: int | None = i + 1
            else:
                pid = None
            groups.setdefault(pid, []).append(d)

        findings: list[ComplianceFinding] = []
        for pid in [*range(1, len(people) + 1), None]:
            items = groups.get(pid, [])
            for rule in self.rules:
                ppe = [d for d in items if d.cls in rule.any_of]
                if rule.body_part == PERSON:
                    if pid is None:
                        continue          # a vest that belongs to nobody: nothing to assess
                    parts = [people[pid - 1]]
                else:
                    parts = [d for d in items if d.cls == rule.body_part]

                outcome = self._judge(rule, parts, ppe, scoped=pid is not None)
                if outcome is None:
                    continue              # neither body part nor PPE seen: not assessed
                status, detail, conf = outcome
                findings.append(
                    ComplianceFinding(
                        body_part=rule.body_part,
                        expected_ppe=" or ".join(rule.any_of),
                        status=status,
                        detail=detail,
                        item=rule.item,
                        level=rule.level,
                        required=rule.level == "required",
                        person_id=pid,
                        confidence=conf,
                    )
                )
        return findings

    def _judge(self, rule: PpeRule, parts: list[Detection], ppe: list[Detection],
               scoped: bool):
        """
        Decide one rule for one person.
        `parts`  body-part boxes of this person (for the vest rule: the person box)
        `ppe`    PPE boxes of this person that satisfy the rule
        `scoped` False when the boxes could not be tied to a single person
        Returns (status, detail, confidence) or None when there is nothing to assess.
        """
        if not parts and not ppe:
            return None
        if not parts and not scoped:
            # A PPE box that belongs to nobody (gloves on a bench, a stray
            # low-confidence box) says nothing about anyone's compliance.
            return None

        floor = self.compliance_conf
        label = " or ".join(rule.any_of)
        part_name = rule.body_part
        on_person = part_name == PERSON   # vest: being on the person IS the match

        def match(part: Detection, p: Detection) -> float:
            return 1.0 if on_person else _iou(part.box_xywhn, p.box_xywhn)

        strong = [p for p in ppe if p.confidence >= floor]
        weak = [p for p in ppe if p.confidence < floor]
        where = "on person" if on_person else f"over {part_name}"

        outcomes: list[tuple[str, str, float]] = []
        for part in parts:
            s_iou, s_box = max(((match(part, p), p) for p in strong),
                               key=lambda t: t[0], default=(0.0, None))
            w_iou, w_box = max(((match(part, p), p) for p in weak),
                               key=lambda t: t[0], default=(0.0, None))
            if s_box is not None and s_iou >= MATCH_IOU:
                outcomes.append((
                    "present",
                    f"{s_box.cls} ({s_box.confidence:.2f}) {where}",
                    s_box.confidence,
                ))
            elif w_box is not None and w_iou >= MATCH_IOU:
                outcomes.append((
                    "uncertain",
                    f"low-confidence {w_box.cls} ({w_box.confidence:.2f}) {where}",
                    w_box.confidence,
                ))
            elif (strong and scoped and len(parts) == 1
                  and part_name not in PAIRED_PARTS):
                outcomes.append((
                    "uncertain",
                    f"{label} detected on this person but not clearly over the "
                    f"{part_name} (max IoU {s_iou:.2f})",
                    max(p.confidence for p in strong),
                ))
            elif part.confidence < floor:
                outcomes.append((
                    "uncertain",
                    f"low-confidence {part_name} detection ({part.confidence:.2f}), "
                    f"no {label} found",
                    part.confidence,
                ))
            else:
                outcomes.append((
                    "missing",
                    f"{part_name} detected ({part.confidence:.2f}) with no {label}",
                    part.confidence,
                ))

        # PPE on the person that sits on no detected body part. A gloved hand is
        # usually not detected as "hands" at all, so the glove box is the evidence.
        unmatched = [p for p in ppe
                     if not any(match(part, p) >= MATCH_IOU for part in parts)]
        for p in unmatched:
            if p.confidence >= floor:
                outcomes.append((
                    "present", f"{p.cls} detected ({p.confidence:.2f})", p.confidence
                ))
        if not outcomes:
            # Only low-confidence PPE boxes and no body part.
            p = max(weak, key=lambda d: d.confidence)
            outcomes.append((
                "uncertain",
                f"low-confidence {p.cls} detection ({p.confidence:.2f})",
                p.confidence,
            ))

        # Worst case wins: one bare hand outweighs one gloved hand.
        worst = max(outcomes, key=lambda o: _RANK[o[0]])
        details: list[str] = []
        for _, text, _ in sorted(outcomes, key=lambda o: -_RANK[o[0]]):
            if text not in details:
                details.append(text)
        return worst[0], "; ".join(details), round(float(worst[2]), 3)