"""Per-frame door-mask loading and YOLO/RTMW-only 2-D risk scoring."""

from __future__ import annotations

import json
import math
import re
import time
from collections import Counter
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np


RASTER_SUFFIXES = {".png", ".bmp", ".tif", ".tiff"}
DEFAULT_DOOR_LABELS = ("door", "车门")


@dataclass
class LoadedDoorMask:
    status: str
    annotation_index: int
    path: Path | None
    mask: np.ndarray | None
    source_format: str | None

    def metadata(self) -> dict:
        return {
            "status": self.status,
            "annotation_index": self.annotation_index,
            "path": str(self.path.resolve()) if self.path is not None else None,
            "format": self.source_format,
            "area_px": int(np.count_nonzero(self.mask)) if self.mask is not None else 0,
        }


def _frame_index_from_path(path: Path) -> int | None:
    matches = re.findall(r"\d+", path.stem)
    return int(matches[-1]) if matches else None


class DoorMaskSequence:
    """Index a directory of X-AnyLabeling JSON polygons or raster masks."""

    def __init__(
        self,
        directory: Path | None = None,
        mask_format: str = "auto",
        labels: tuple[str, ...] = DEFAULT_DOOR_LABELS,
        frame_offset: int = 0,
        require_every_frame: bool = False,
        constant_mask: Path | None = None,
    ):
        if (directory is None) == (constant_mask is None):
            raise ValueError("Provide exactly one of directory or constant_mask")
        self.constant_path = (
            Path(constant_mask).expanduser().resolve()
            if constant_mask is not None
            else None
        )
        self.directory = (
            self.constant_path.parent if self.constant_path is not None else Path(directory)
        )
        if self.constant_path is not None and not self.constant_path.is_file():
            raise ValueError(f"Constant door mask does not exist: {self.constant_path}")
        if self.constant_path is None and not self.directory.is_dir():
            raise ValueError(f"Door mask directory does not exist: {self.directory}")
        self.labels = {label.casefold() for label in labels}
        self.frame_offset = int(frame_offset)
        self.require_every_frame = require_every_frame

        if self.constant_path is not None:
            if mask_format == "auto":
                selected_format = (
                    "xanylabeling"
                    if self.constant_path.suffix.lower() == ".json"
                    else "raster"
                )
            else:
                selected_format = mask_format
            self.source_format = selected_format
            self.paths = {0: self.constant_path}
            return

        json_paths = sorted(self.directory.rglob("*.json"))
        raster_paths = sorted(
            path
            for path in self.directory.rglob("*")
            if path.is_file() and path.suffix.lower() in RASTER_SUFFIXES
        )
        if mask_format == "auto":
            selected_format = "xanylabeling" if json_paths else "raster"
        else:
            selected_format = mask_format
        candidates = json_paths if selected_format == "xanylabeling" else raster_paths
        if not candidates:
            expected = "JSON" if selected_format == "xanylabeling" else "PNG/BMP/TIFF"
            raise ValueError(f"No {expected} door masks found in {self.directory}")

        self.source_format = selected_format
        self.paths: dict[int, Path] = {}
        for path in candidates:
            index = _frame_index_from_path(path)
            if index is None and selected_format == "xanylabeling":
                data = json.loads(path.read_text(encoding="utf-8-sig"))
                explicit = data.get("frame_index", data.get("frameIndex"))
                index = int(explicit) if explicit is not None else None
            if index is None:
                continue
            if index in self.paths:
                raise ValueError(
                    f"Multiple door masks map to annotation frame {index}: "
                    f"{self.paths[index]} and {path}"
                )
            self.paths[index] = path
        if not self.paths:
            raise ValueError(
                "Door mask filenames need a frame number, for example 000000.json"
            )

    def load(self, video_frame_index: int, frame_shape: tuple[int, int]) -> LoadedDoorMask:
        if self.constant_path is not None:
            annotation_index = video_frame_index
            path = self.constant_path
        else:
            annotation_index = video_frame_index + self.frame_offset
            path = self.paths.get(annotation_index)
        if path is None:
            if self.require_every_frame:
                raise RuntimeError(
                    f"No door mask for video frame {video_frame_index} "
                    f"(expected annotation index {annotation_index})"
                )
            return LoadedDoorMask("missing", annotation_index, None, None, self.source_format)

        if self.source_format == "xanylabeling":
            mask = self._load_xanylabeling(path, frame_shape)
        else:
            mask = self._load_raster(path, frame_shape)
        status = "loaded" if np.any(mask) else "empty"
        return LoadedDoorMask(status, annotation_index, path, mask, self.source_format)

    @staticmethod
    def _load_raster(path: Path, frame_shape: tuple[int, int]) -> np.ndarray:
        image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if image is None:
            raise RuntimeError(f"Cannot read door mask: {path}")
        if image.ndim == 3:
            image = np.max(image[:, :, :3], axis=2)
        height, width = frame_shape
        if image.shape != (height, width):
            image = cv2.resize(image, (width, height), interpolation=cv2.INTER_NEAREST)
        return image > 0

    def _load_xanylabeling(
        self, path: Path, frame_shape: tuple[int, int]
    ) -> np.ndarray:
        data = json.loads(path.read_text(encoding="utf-8-sig"))
        height, width = frame_shape
        source_width = float(data.get("imageWidth") or width)
        source_height = float(data.get("imageHeight") or height)
        scale = np.asarray(
            [width / max(source_width, 1.0), height / max(source_height, 1.0)],
            dtype=np.float32,
        )
        mask = np.zeros((height, width), dtype=np.uint8)
        for shape in data.get("shapes", []):
            if str(shape.get("label", "")).casefold() not in self.labels:
                continue
            points = np.asarray(shape.get("points", []), dtype=np.float32).reshape(-1, 2)
            if len(points) < 2:
                continue
            points = np.rint(points * scale).astype(np.int32)
            shape_type = str(shape.get("shape_type", "polygon")).lower()
            if shape_type == "rectangle" and len(points) >= 2:
                x1, y1 = points[0]
                x2, y2 = points[1]
                cv2.rectangle(mask, (x1, y1), (x2, y2), 1, -1)
            elif shape_type == "circle" and len(points) >= 2:
                radius = int(round(float(np.linalg.norm(points[1] - points[0]))))
                cv2.circle(mask, tuple(points[0]), radius, 1, -1)
            elif len(points) >= 3:
                cv2.fillPoly(mask, [points], 1)
        return mask > 0


def _fit_mask(mask: np.ndarray, frame_shape: tuple[int, int]) -> np.ndarray:
    height, width = frame_shape
    result = np.asarray(mask) > 0
    if result.shape != (height, width):
        result = cv2.resize(
            result.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST
        ) > 0
    return result


def _box_mask(box: list[float] | None, frame_shape: tuple[int, int]) -> np.ndarray | None:
    if box is None or len(box) < 4:
        return None
    height, width = frame_shape
    x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]]
    x1, x2 = sorted((max(0, min(width - 1, x1)), max(0, min(width - 1, x2))))
    y1, y2 = sorted((max(0, min(height - 1, y1)), max(0, min(height - 1, y2))))
    if x2 <= x1 or y2 <= y1:
        return None
    mask = np.zeros((height, width), dtype=bool)
    mask[y1 : y2 + 1, x1 : x2 + 1] = True
    return mask


def _hand_mask(points: list, frame_shape: tuple[int, int]) -> np.ndarray | None:
    hand = np.asarray(points, dtype=np.float32).reshape(-1, 2)
    if len(hand) < 3 or not np.isfinite(hand).all():
        return None
    hull = cv2.convexHull(np.rint(hand).astype(np.int32))
    mask = np.zeros(frame_shape, dtype=np.uint8)
    cv2.fillConvexPoly(mask, hull, 1)
    return mask > 0 if np.any(mask) else None


def _decision(distance_relative: float, overlap_px: int, warning: float, release: float) -> tuple[str, str]:
    if overlap_px > 0:
        return "danger", "DANGER_INTERSECTION"
    if distance_relative <= warning:
        return "warning", "WARNING_NEAR"
    if distance_relative <= release:
        return "watch", "WATCH"
    return "clear", "CLEAR"


class DoorRiskEvaluator:
    """Score RTMW hand regions, with one YOLO body-box fallback per person."""

    def __init__(
        self,
        sequence: DoorMaskSequence,
        warning_relative: float = 0.30,
        release_relative: float = 0.45,
    ):
        if warning_relative < 0 or release_relative <= warning_relative:
            raise ValueError("Risk thresholds require 0 <= warning < release")
        self.sequence = sequence
        self.warning_relative = float(warning_relative)
        self.release_relative = float(release_relative)
        self.mask_status_counts: Counter = Counter()
        self.level_counts: Counter = Counter()
        self.source_counts: Counter = Counter()
        self.maximum_score = 0.0
        self.timing_seconds: Counter = Counter()
        self.timing_counts: Counter = Counter()
        self._constant_cache_shape: tuple[int, int] | None = None
        self._constant_loaded: LoadedDoorMask | None = None
        self._constant_distance_map: np.ndarray | None = None
        self.distance_map_cache_builds = 0

    def _record_timing(self, name: str, started: float) -> None:
        self.timing_seconds[name] += time.perf_counter() - started
        self.timing_counts[name] += 1

    def _measure_region(
        self,
        region: np.ndarray,
        door_mask: np.ndarray,
        distance_map: np.ndarray,
        source: str,
        track_id: int,
        side: str,
        state: str | None,
    ) -> dict:
        area_px = int(np.count_nonzero(region))
        overlap_px = int(np.count_nonzero(region & door_mask))
        distances = distance_map[region]
        distance_px = float(np.min(distances)) if len(distances) else math.inf
        scale_px = max(math.sqrt(area_px), 1.0)
        distance_relative = distance_px / scale_px
        score = float(np.clip(1.0 - distance_relative / self.release_relative, 0.0, 1.0))
        if overlap_px:
            score = 1.0
        level, decision = _decision(
            distance_relative,
            overlap_px,
            self.warning_relative,
            self.release_relative,
        )
        result = {
            "status": "calculated",
            "risk_score": round(score, 5),
            "risk_level": level,
            "decision": decision,
            "source": source,
            "track_id": int(track_id),
            "side": side,
            "state": state,
            "region_area_px": area_px,
            "door_overlap_px": overlap_px,
            "region_overlap_ratio": round(overlap_px / max(area_px, 1), 6),
            "min_distance_px": round(distance_px, 3),
            "region_scale_px": round(scale_px, 3),
            "distance_relative": round(distance_relative, 6),
        }
        self.level_counts[level] += 1
        self.source_counts[source] += 1
        self.maximum_score = max(self.maximum_score, score)
        return result

    @staticmethod
    def _unavailable(source: str, track_id: int, side: str, status: str) -> dict:
        return {
            "status": status,
            "risk_score": None,
            "risk_level": "unavailable",
            "decision": "EVIDENCE_UNAVAILABLE",
            "source": source,
            "track_id": int(track_id),
            "side": side,
        }

    def evaluate(
        self,
        frame_index: int,
        record: dict,
        frame_shape: tuple[int, int],
    ) -> tuple[dict, np.ndarray | None]:
        total_started = time.perf_counter()
        stage_started = time.perf_counter()
        use_constant_cache = self.sequence.constant_path is not None
        if use_constant_cache and self._constant_cache_shape == frame_shape:
            cached = self._constant_loaded
            assert cached is not None
            loaded = LoadedDoorMask(
                cached.status,
                frame_index,
                cached.path,
                cached.mask,
                cached.source_format,
            )
        else:
            loaded = self.sequence.load(frame_index, frame_shape)
            if use_constant_cache:
                self._constant_cache_shape = frame_shape
                self._constant_loaded = loaded
                self._constant_distance_map = None
        self._record_timing("door_mask_load", stage_started)
        self.mask_status_counts[loaded.status] += 1
        mask_ready = loaded.status == "loaded" and loaded.mask is not None
        door_mask = loaded.mask if mask_ready else None
        distance_map = None
        if door_mask is not None and use_constant_cache and self._constant_distance_map is not None:
            distance_map = self._constant_distance_map
        elif door_mask is not None:
            stage_started = time.perf_counter()
            distance_map = cv2.distanceTransform(
                (~door_mask).astype(np.uint8), cv2.DIST_L2, cv2.DIST_MASK_PRECISE
            )
            self._record_timing("distance_transform", stage_started)
            if use_constant_cache:
                self._constant_distance_map = distance_map
                self.distance_map_cache_builds += 1

        stage_started = time.perf_counter()
        regions: list[tuple[dict, np.ndarray, str, int, str, str | None]] = []
        track_ids = [int(value) for value in record.get("track_ids", [])]
        person_positions_with_hands: set[int] = set()
        for hand in record.get("global_hands", []):
            position = int(hand.get("person_id", 0)) - 1
            track_id = track_ids[position] if 0 <= position < len(track_ids) else position + 1
            side = str(hand.get("side", "unknown"))
            region = _hand_mask(hand.get("keypoints", []), frame_shape)
            if region is None:
                continue
            person_positions_with_hands.add(position)
            regions.append(
                (
                    hand,
                    region,
                    "RTMW_HAND_KEYPOINTS",
                    track_id,
                    side,
                    str(hand.get("state", "unknown")),
                )
            )

        boxes = record.get("boxes", [])
        body_fallbacks = []
        for position, box in enumerate(boxes):
            if position in person_positions_with_hands:
                continue
            region = _box_mask(box, frame_shape)
            if region is None:
                continue
            track_id = track_ids[position] if position < len(track_ids) else position + 1
            fallback = {
                "track_id": track_id,
                "side": "body",
                "state": "rtmw_hand_unavailable",
                "body_zone_box": box,
            }
            body_fallbacks.append(fallback)
            regions.append(
                (
                    fallback,
                    region,
                    "YOLO_BODY_BOX_FALLBACK",
                    track_id,
                    "body",
                    "rtmw_hand_unavailable",
                )
            )
        record["body_fallbacks"] = body_fallbacks
        self._record_timing("region_build", stage_started)

        stage_started = time.perf_counter()
        measurements = []
        unavailable_status = f"door_mask_{loaded.status}"
        for owner, region, source, track_id, side, state in regions:
            if door_mask is None or distance_map is None:
                measurement = self._unavailable(
                    source, track_id, side, unavailable_status
                )
            else:
                measurement = self._measure_region(
                    region, door_mask, distance_map, source, track_id, side, state
                )
            owner["risk"] = measurement
            measurements.append(measurement)
        self._record_timing("region_measure", stage_started)

        calculated = [item for item in measurements if item["risk_score"] is not None]
        nearest = max(calculated, key=lambda item: item["risk_score"]) if calculated else None
        if loaded.status != "loaded":
            status = unavailable_status
        elif not calculated:
            status = "person_region_missing"
        else:
            status = "calculated"
        frame_risk = {
            "status": status,
            "risk_score": nearest["risk_score"] if nearest else None,
            "risk_level": nearest["risk_level"] if nearest else "unavailable",
            "decision": nearest["decision"] if nearest else "EVIDENCE_UNAVAILABLE",
            "source": nearest["source"] if nearest else None,
            "track_id": nearest["track_id"] if nearest else None,
            "side": nearest["side"] if nearest else None,
            "door_mask": loaded.metadata(),
            "warning_distance_relative": self.warning_relative,
            "release_distance_relative": self.release_relative,
            "regions": measurements,
        }
        self._record_timing("risk_total", total_started)
        return frame_risk, loaded.mask

    def summary(self) -> dict:
        return {
            "mask_directory": str(self.sequence.directory.resolve()),
            "constant_mask": (
                str(self.sequence.constant_path)
                if self.sequence.constant_path is not None
                else None
            ),
            "constant_for_all_frames": self.sequence.constant_path is not None,
            "mask_format": self.sequence.source_format,
            "door_labels": sorted(self.sequence.labels),
            "frame_offset": self.sequence.frame_offset,
            "warning_distance_relative": self.warning_relative,
            "release_distance_relative": self.release_relative,
            "mask_status_counts": dict(self.mask_status_counts),
            "risk_level_counts": dict(self.level_counts),
            "source_counts": dict(self.source_counts),
            "maximum_risk_score": round(self.maximum_score, 5),
            "constant_distance_map_cached": self.sequence.constant_path is not None,
            "distance_map_cache_builds": self.distance_map_cache_builds,
            "timing_ms": {
                name: {
                    "count": int(self.timing_counts[name]),
                    "total": round(float(seconds) * 1000.0, 3),
                    "mean": round(
                        float(seconds) * 1000.0
                        / max(int(self.timing_counts[name]), 1),
                        3,
                    ),
                }
                for name, seconds in self.timing_seconds.items()
            },
            "score_semantics": "2-D image overlap/proximity research score; not millimeters or safety certification",
            "region_policy": "RTMW alarm-eligible hand hulls; YOLO body box only when a person has no usable RTMW hand region",
        }