"""Strict consumer for an Ultralytics end-to-end-NMS ONNX (person element).

Contract: static batch-1 NCHW input, one ``[1, 300, 6]`` output whose rows are
``x1, y1, x2, y2, confidence, model_class_id``.  NMS runs inside the graph, so
this module only validates the tensors, applies the confidence gate, undoes the
letterbox, drops geometrically impossible boxes and caps the result count.

No ``dataclass`` declarations: ScoreVision's ``spec_from_file_location`` loader
does not register the module in ``sys.modules`` before executing it, which makes
dataclass decoration fail at import.

ELEMENT_CONFIG below is written by ``scripts/build_submission.py`` from the
policy chosen on the validation split -- do not hand-edit it without re-running
``scripts/optimize_policy.py``, since ``nms_iou`` in particular must match the
value frozen into the ONNX graph at export time.
"""

import ast
import json
import math
import os
from pathlib import Path
from typing import Any

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

API_VERSION = "c2-strict-v1"
STRICT_INFERENCE_ENV = "SV_STRICT_INFERENCE"
OUTPUT_ROWS = 300
OUTPUT_COLUMNS = 6

ELEMENT_CONFIG: dict[str, Any] = {
    "element_id": "person",
    "classes": ["person"],
    "coco_ids": [0],
    "imgsz": 704,
    "conf": 0.55,
    "nms_iou": 0.45,  # frozen into the ONNX graph at export
    "max_det": 300,
    "min_side": 3,
    "max_aspect": 10,
    "max_area_frac": 0.95,
    "scale_w": 1,  # box geometry calibration to the validator standard
    "scale_h": 1,
    "conf_steps": ((30, 0.55), (60, 0.6), (None, 0.65)),
}


class InferenceInvariantError(RuntimeError):
    """A model tensor or decoded detection violated the serving contract."""


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox]
    keypoints: list[tuple[int, int]]


def _conf_for(n_raw: int) -> float:
    """Confidence gate, optionally raised on crowded frames.

    The wrong-box penalty is a per-image absolute count that saturates at ten,
    so a frame holding ninety people can lose the whole term while a frame with
    six cannot.  When CONF_STEPS is non-empty the gate is spent where that
    exposure actually is.  Empty steps means a flat gate.
    """
    for limit, conf in ELEMENT_CONFIG.get("conf_steps", ()):
        if limit is None or n_raw < limit:
            return float(conf)
    return float(ELEMENT_CONFIG["conf"])


class Miner:
    api_version = API_VERSION
    strict_inference = False

    def __init__(self, path_hf_repo: Path) -> None:
        cfg = ELEMENT_CONFIG
        self.config = cfg
        self.class_names = list(cfg["classes"])
        self.conf_threshold = float(cfg["conf"])
        self.max_det = int(cfg["max_det"])
        self.min_side = float(cfg.get("min_side", 3.0))
        self.max_aspect_ratio = float(cfg.get("max_aspect", 10.0))
        self.max_area_frac = float(cfg.get("max_area_frac", 0.95))
        self.conf_steps = tuple(tuple(s) for s in cfg.get("conf_steps", ()))
        # box geometry calibration: the validator draws wider/taller boxes than
        # the training sources; scale about the centre before the IoU match.
        self.scale_w = float(cfg.get("scale_w", 1.0))
        self.scale_h = float(cfg.get("scale_h", 1.0))
        raw = os.environ.get(STRICT_INFERENCE_ENV, "")
        if raw not in ("", "0", "1"):
            raise ValueError(f"{STRICT_INFERENCE_ENV} must be '1' or unset/'0', got {raw!r}")
        self.strict_inference = raw == "1"
        self.inference_failures = 0

        # The referee grants one physical core (two SMT threads).  Let ORT own
        # both for the convolution graph and keep OpenCV single-threaded: a
        # second OpenCV worker is slower and adds tail jitter on two siblings.
        cv2.setNumThreads(1)
        options = ort.SessionOptions()
        options.intra_op_num_threads = 2
        options.inter_op_num_threads = 1
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        model_path = Path(path_hf_repo) / "weights.onnx"
        try:
            self.session = ort.InferenceSession(
                str(model_path), sess_options=options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
        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"]
            )

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        shape = self.session.get_inputs()[0].shape
        if len(shape) != 4 or shape[1] != 3:
            raise ValueError(f"model input must be rank-4 NCHW with 3 channels, got {shape!r}")
        if not all(isinstance(v, int) and v > 0 for v in shape[2:]):
            raise ValueError(f"e2e export must have static spatial dims, got {shape[2:]!r}")
        if (shape[2], shape[3]) != (cfg["imgsz"], cfg["imgsz"]):
            raise ValueError(f"model input {shape[2]}x{shape[3]} != config imgsz {cfg['imgsz']}")
        self.input_height = self.input_width = int(cfg["imgsz"])
        out_shape = self.session.get_outputs()[0].shape
        if list(out_shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
            raise ValueError(f"e2e output must be [1, {OUTPUT_ROWS}, {OUTPUT_COLUMNS}], got {out_shape!r}")
        self.model_names = self._read_model_names()
        self._warmup(3)
        print(
            f"api_version={self.api_version} element={cfg['element_id']} imgsz={cfg['imgsz']} "
            f"conf={self.conf_threshold} steps={self.conf_steps} "
            f"providers={self.session.get_providers()}"
        )

    def _read_model_names(self) -> tuple[str, ...]:
        try:
            raw = self.session.get_modelmeta().custom_metadata_map["names"]
            parsed = ast.literal_eval(raw)
            if not isinstance(parsed, dict):
                raise TypeError("names metadata is not a dict")
            keys = sorted(parsed)
            if keys != list(range(len(keys))):
                raise ValueError(f"names metadata keys not consecutive from zero: {keys[:8]}")
            return tuple(str(parsed[i]) for i in keys)
        except Exception as exc:
            raise ValueError(f"e2e ONNX requires valid names metadata: {exc}") from exc

    @staticmethod
    def _letterbox(image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114)):
        height, width = image.shape[:2]
        new_width, new_height = new_shape
        ratio = min(new_width / width, new_height / height)
        rw, rh = int(round(width * ratio)), int(round(height * ratio))
        if (rw, rh) != (width, height):
            image = cv2.resize(
                image, (rw, rh),
                interpolation=cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR,
            )
        dw, dh = (new_width - rw) / 2.0, (new_height - rh) / 2.0
        padded = cv2.copyMakeBorder(
            image,
            int(round(dh - 0.1)), int(round(dh + 0.1)),
            int(round(dw - 0.1)), int(round(dw + 0.1)),
            cv2.BORDER_CONSTANT, value=color,
        )
        return padded, ratio, (dw, dh)

    def _decode(self, output, ratio, pad, original_size):
        rows = output[0]
        boxes = rows[:, :4].astype(np.float32, copy=True)
        scores = rows[:, 4].astype(np.float32, copy=False)
        raw_classes = rows[:, 5]
        if (scores < 0).any() or (scores > 1).any():
            raise InferenceInvariantError("e2e confidence outside [0, 1]")
        rounded = np.rint(raw_classes)
        if not np.array_equal(raw_classes, rounded):
            raise InferenceInvariantError("e2e model class ids must be exact integers")
        model_classes = rounded.astype(np.int64)
        if (model_classes < 0).any() or (model_classes >= len(self.model_names)).any():
            raise InferenceInvariantError("e2e model class id outside range")

        active = scores > 0.0
        reversed_rows = active & ((boxes[:, 2] < boxes[:, 0]) | (boxes[:, 3] < boxes[:, 1]))
        if reversed_rows.any():
            raise InferenceInvariantError("e2e output contains reversed boxes")

        # A floor below every reachable gate; nothing under it can be emitted.
        floor = min([self.conf_threshold] + [float(c) for _, c in self.conf_steps])
        keep = active & (model_classes == 0) & (scores >= np.float32(floor))
        if not keep.any():
            return np.empty((0, 4), np.float32), np.empty((0,), np.float32)
        boxes, scores = boxes[keep], scores[keep]

        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad[0]) / ratio
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad[1]) / ratio
        width, height = original_size
        if self.scale_w != 1.0 or self.scale_h != 1.0:
            cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
            cy = (boxes[:, 1] + boxes[:, 3]) * 0.5
            hw = (boxes[:, 2] - boxes[:, 0]) * 0.5 * np.float32(self.scale_w)
            hh = (boxes[:, 3] - boxes[:, 1]) * 0.5 * np.float32(self.scale_h)
            boxes = np.stack([cx - hw, cy - hh, cx + hw, cy + hh], 1).astype(np.float32)
        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width - 1)
        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height - 1)

        bw = boxes[:, 2] - boxes[:, 0]
        bh = boxes[:, 3] - boxes[:, 1]
        aspect = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
        sane = (
            (bw >= self.min_side)
            & (bh >= self.min_side)
            & (bw * bh <= self.max_area_frac * float(width * height))
            & (aspect <= self.max_aspect_ratio)
        )
        boxes, scores = boxes[sane], scores[sane]
        if not len(scores):
            return boxes, scores
        order = np.argsort(-scores, kind="stable")[: self.max_det]
        boxes, scores = boxes[order], scores[order]
        gate = _conf_for(len(scores)) if self.conf_steps else self.conf_threshold
        emit = scores >= np.float32(gate)
        return boxes[emit], scores[emit]

    def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
        if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
            raise ValueError(f"expected HWC 3-channel image, got {getattr(image, 'shape', type(image))}")
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)
        original_height, original_width = image.shape[:2]
        padded, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
        blob = cv2.dnn.blobFromImage(padded, scalefactor=1 / 255.0, swapRB=True)
        outputs = self.session.run(self.output_names, {self.input_name: blob})
        output = outputs[0]
        if not np.isfinite(output).all():
            raise InferenceInvariantError("model output must be finite everywhere")
        if list(output.shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
            raise ValueError(f"unexpected output shape {list(output.shape)}")
        boxes, scores = self._decode(output, ratio, pad, (original_width, original_height))
        return [
            BoundingBox(
                x1=int(math.floor(b[0])), y1=int(math.floor(b[1])),
                x2=int(math.ceil(b[2])), y2=int(math.ceil(b[3])),
                cls_id=0, conf=float(s),
            )
            for b, s in zip(boxes, scores)
            if b[2] > b[0] and b[3] > b[1]
        ]

    def _warmup(self, iterations: int) -> None:
        dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
        original = self.strict_inference
        self.strict_inference = True
        try:
            for _ in range(iterations):
                self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
        except Exception as exc:
            raise RuntimeError(f"warmup failed, refusing to serve: {type(exc).__name__}: {exc}") from exc
        finally:
            self.strict_inference = original

    def predict_batch(self, batch_images, offset: int, n_keypoints: int) -> list[TVFrameResult]:
        _ = n_keypoints
        results: list[TVFrameResult] = []
        for index, image in enumerate(batch_images):
            frame_id = offset + index
            try:
                boxes = self._predict_single(image)
            except Exception as exc:
                if self.strict_inference:
                    raise
                self.inference_failures += 1
                print(json.dumps({
                    "event": "inference_failure",
                    "frame_id": frame_id,
                    "error_type": type(exc).__name__,
                    "error": str(exc),
                    "inference_failures": self.inference_failures,
                }, sort_keys=True))
                boxes = []
            results.append(TVFrameResult(frame_id=frame_id, boxes=boxes, keypoints=[]))
        return results

    def __repr__(self) -> str:
        return f"E2EONNXRuntime(providers={self.session.get_providers()})"