"""Vehicle stage: a COCO object detector beside the plate recognizer (Job 2, phase A).

Answers "what KIND of vehicle is in this entry frame?" for the Car Wash desk's category
suggestion (wiki/decisions/venue-modules.md §Vehicle category from vision). ADVISORY by
design: the Node server records it next to the plate, the desk pre-selects the site
category it maps to, the operator decides, a confident downgrade is flagged. Nothing is
ever gated on it, so a wrong or missing detection costs nothing but a suggestion.

Model: YOLOX (Megvii, Apache-2.0) as an ONNX graph on the ONNX Runtime the plate stage
already uses — the licence rule that keeps Ultralytics (AGPL) out. COCO's vehicle classes
are car / motorcycle / bus / truck: enough to tell a van or a truck from a car, NOT enough
for SUV vs sedan — that is phase B (a body-type classifier on the pilot's own frames).
The detector's vehicle box is also the crop phase B will classify.

Pure numpy/cv2 pre/post-processing, no torch: letterbox to the model's square input
(pad 114, no normalisation — YOLOX's exported graphs take raw 0–255 BGR), decode the
stride grids, class-agnostic NMS, map COCO ids to the shared vocabulary, pick ONE
vehicle: the one whose box holds the plate (when a plate was read), else the largest.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

from .schemas import BBox, VehicleResult

# COCO-80 class index → the shared VEHICLE_CLASSES vocabulary (packages/shared).
COCO_VEHICLE_CLASSES: dict[int, str] = {2: "car", 3: "motorcycle", 5: "bus", 7: "truck"}

# YOLOX feature strides; grids are input/stride per level (8400 anchors at 640).
_STRIDES = (8, 16, 32)


@dataclass(frozen=True)
class Detection:
    body_type: str
    confidence: float
    x1: float
    y1: float
    x2: float
    y2: float

    @property
    def area(self) -> float:
        return max(0.0, self.x2 - self.x1) * max(0.0, self.y2 - self.y1)

    def contains(self, x: float, y: float) -> bool:
        return self.x1 <= x <= self.x2 and self.y1 <= y <= self.y2


class VehicleDetector(Protocol):
    """What the recognizer composition needs: frame bytes (+ the plate box) → a class."""

    @property
    def model_version(self) -> str: ...

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None: ...


# ----------------------------------------------------------------------------------
# Pre/post-processing (pure functions — unit-tested on synthetic tensors)
# ----------------------------------------------------------------------------------


def letterbox(frame: Any, size: int) -> tuple[Any, float]:
    """Resize keeping aspect, pad bottom/right with 114 to size×size. Returns the CHW
    float32 tensor (batch dim added) and the scale to map boxes back."""
    import cv2
    import numpy as np

    h, w = frame.shape[:2]
    r = min(size / h, size / w)
    nh, nw = int(round(h * r)), int(round(w * r))
    resized = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_LINEAR)
    padded = np.full((size, size, 3), 114, dtype=np.uint8)
    padded[:nh, :nw] = resized
    tensor = padded.transpose(2, 0, 1)[None].astype(np.float32)
    return np.ascontiguousarray(tensor), r


def decode(raw: Any, size: int) -> Any:
    """YOLOX raw output [N, 5+classes] (batch squeezed) → same shape with xywh decoded
    into pixel units of the letterboxed input. Rows are ordered stride 8, 16, 32."""
    import numpy as np

    out = raw.astype(np.float32).copy()
    grids = []
    strides = []
    for s in _STRIDES:
        n = size // s
        ys, xs = np.meshgrid(np.arange(n), np.arange(n), indexing="ij")
        grids.append(np.stack((xs, ys), axis=-1).reshape(-1, 2))
        strides.append(np.full((n * n, 1), s, dtype=np.float32))
    grid = np.concatenate(grids, axis=0).astype(np.float32)
    stride = np.concatenate(strides, axis=0)
    if out.shape[0] != grid.shape[0]:
        raise ValueError(f"unexpected output rows {out.shape[0]} for input {size} (want {grid.shape[0]})")
    out[:, :2] = (out[:, :2] + grid) * stride
    out[:, 2:4] = np.exp(out[:, 2:4]) * stride
    return out


def nms(boxes: Any, scores: Any, iou_threshold: float) -> list[int]:
    """Greedy class-agnostic non-max suppression over xyxy boxes; returns kept indices."""
    import numpy as np

    if len(boxes) == 0:
        return []
    order = scores.argsort()[::-1]
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
    keep: list[int] = []
    while order.size > 0:
        i = int(order[0])
        keep.append(i)
        if order.size == 1:
            break
        rest = order[1:]
        xx1 = np.maximum(x1[i], x1[rest])
        yy1 = np.maximum(y1[i], y1[rest])
        xx2 = np.minimum(x2[i], x2[rest])
        yy2 = np.minimum(y2[i], y2[rest])
        inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None)
        iou = inter / (areas[i] + areas[rest] - inter + 1e-9)
        order = rest[iou <= iou_threshold]
    return keep


def vehicles_from_output(
    raw: Any, size: int, scale: float, min_confidence: float, iou_threshold: float = 0.45
) -> list[Detection]:
    """Full post-processing: decode → vehicle classes only → confidence floor → NMS →
    boxes in ORIGINAL frame pixels."""
    import numpy as np

    dec = decode(raw, size)
    cls_scores = dec[:, 5:]
    cls_idx = cls_scores.argmax(axis=1)
    score = dec[:, 4] * cls_scores[np.arange(len(dec)), cls_idx]
    wanted = np.isin(cls_idx, list(COCO_VEHICLE_CLASSES)) & (score >= min_confidence)
    if not wanted.any():
        return []
    d = dec[wanted]
    s = score[wanted]
    c = cls_idx[wanted]
    boxes = np.stack(
        (d[:, 0] - d[:, 2] / 2, d[:, 1] - d[:, 3] / 2, d[:, 0] + d[:, 2] / 2, d[:, 1] + d[:, 3] / 2), axis=1
    )
    keep = nms(boxes, s, iou_threshold)
    out: list[Detection] = []
    for i in keep:
        b = boxes[i] / scale
        out.append(
            Detection(
                body_type=COCO_VEHICLE_CLASSES[int(c[i])],
                confidence=float(s[i]),
                x1=float(b[0]),
                y1=float(b[1]),
                x2=float(b[2]),
                y2=float(b[3]),
            )
        )
    return out


def pick_vehicle(detections: list[Detection], plate: BBox | None) -> Detection | None:
    """ONE vehicle per frame: the box holding the plate's centre (the car that was read —
    a lane frame can show the car behind too), else the largest box (nearest the camera)."""
    if not detections:
        return None
    if plate is not None:
        cx = (plate.x1 + plate.x2) / 2
        cy = (plate.y1 + plate.y2) / 2
        holders = [d for d in detections if d.contains(cx, cy)]
        if holders:
            return min(holders, key=lambda d: d.area)  # the tightest box around the plate
    return max(detections, key=lambda d: d.area)


# ----------------------------------------------------------------------------------
# The ONNX Runtime detector
# ----------------------------------------------------------------------------------


class YoloxVehicleDetector:
    """YOLOX ONNX on onnxruntime (CPU). Loads once; a load failure is surfaced through
    `error` and the stage simply yields no vehicle (never breaks the plate path)."""

    def __init__(self, model_path: str, input_size: int = 640, min_confidence: float = 0.4) -> None:
        self._path = Path(model_path)
        self._size = input_size
        self._min_confidence = min_confidence
        self._session = None
        self._input_name = "images"
        self._error: str | None = None
        try:
            import onnxruntime as ort

            opts = ort.SessionOptions()
            opts.intra_op_num_threads = 2  # one frame per entry; leave cores to the lane
            self._session = ort.InferenceSession(
                str(self._path), sess_options=opts, providers=["CPUExecutionProvider"]
            )
            self._input_name = self._session.get_inputs()[0].name
        except Exception as exc:  # noqa: BLE001 - not-ready, never fatal
            self._error = f"{type(exc).__name__}: {exc}"

    @property
    def model_version(self) -> str:
        return f"yolox:{self._path.name}@{self._size}"

    @property
    def ready(self) -> bool:
        return self._session is not None

    @property
    def error(self) -> str | None:
        return self._error

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
        if self._session is None:
            return None
        import cv2
        import numpy as np

        frame = cv2.imdecode(np.frombuffer(image_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        if frame is None:
            return None
        return self.detect_frame(frame, plate)

    def detect_frame(self, frame: Any, plate: BBox | None) -> VehicleResult | None:
        """Same as detect() on an already-decoded BGR frame (the classifier stage decodes
        once and shares it)."""
        if self._session is None:
            return None
        tensor, scale = letterbox(frame, self._size)
        raw = self._session.run(None, {self._input_name: tensor})[0][0]
        found = vehicles_from_output(raw, self._size, scale, self._min_confidence)
        best = pick_vehicle(found, plate)
        if best is None:
            return None
        h, w = frame.shape[:2]
        box = BBox(
            x1=max(0, int(best.x1)), y1=max(0, int(best.y1)), x2=min(w, int(best.x2)), y2=min(h, int(best.y2))
        )
        return VehicleResult(body_type=best.body_type, confidence=round(best.confidence, 4), bbox=box)


# ----------------------------------------------------------------------------------
# Phase B: the body-type classifier on the detector's crop
# ----------------------------------------------------------------------------------

SIDECAR_FORMAT = "parking-bodytype/1"


def crop_vehicle(frame: Any, box: BBox, plate: BBox | None, margin: float) -> Any:
    """The detector's box + margin, plate blurred — the SAME cut the collector stores
    (apps/server review-outbox.ts makeReviewCrop), so the classifier sees at the booth
    what it was trained on. Returns a BGR array, or None when the box is degenerate."""
    import cv2

    h, w = frame.shape[:2]
    mw = round((box.x2 - box.x1) * margin)
    mh = round((box.y2 - box.y1) * margin)
    left, top = max(0, box.x1 - mw), max(0, box.y1 - mh)
    right, bottom = min(w, box.x2 + mw), min(h, box.y2 + mh)
    if right - left < 8 or bottom - top < 8:
        return None
    crop = frame[top:bottom, left:right].copy()
    if plate is not None:
        pad = round(max(plate.x2 - plate.x1, plate.y2 - plate.y1) * 0.25)
        pl, pt = max(0, plate.x1 - pad - left), max(0, plate.y1 - pad - top)
        pr, pb = min(right - left, plate.x2 + pad - left), min(bottom - top, plate.y2 + pad - top)
        if pr - pl >= 2 and pb - pt >= 2:
            sigma = max(6, round((pr - pl) / 6))
            crop[pt:pb, pl:pr] = cv2.GaussianBlur(crop[pt:pb, pl:pr], (0, 0), sigma)
    return crop


class BodyTypeClassifier:
    """`bodytype.onnx` + its `bodytype.json` sidecar (written by apps/trainer). The sidecar
    carries the preprocessing contract — class list, input size, crop margin — and the graph
    normalises internally, so this side only cuts, resizes (INTER_AREA, like the trainer)
    and feeds raw RGB 0–255. Load failure → `error`, the stage yields nothing."""

    def __init__(self, model_path: str, min_confidence: float = 0.6) -> None:
        import json

        self._path = Path(model_path)
        self.min_confidence = min_confidence
        self._session = None
        self._input_name = "image"
        self._error: str | None = None
        self.classes: list[str] = []
        self.version = "?"
        self.input_size = 224
        self.crop_margin = 0.08
        try:
            side = json.loads(self._path.with_suffix(".json").read_text())
            if side.get("format") != SIDECAR_FORMAT:
                raise ValueError(f"unknown sidecar format {side.get('format')!r}")
            self.classes = [str(c) for c in side["classes"]]
            self.version = str(side.get("version", "?"))
            self.input_size = int(side.get("input_size", 224))
            self.crop_margin = float(side.get("crop_margin", 0.08))
            import onnxruntime as ort

            opts = ort.SessionOptions()
            opts.intra_op_num_threads = 2
            self._session = ort.InferenceSession(
                str(self._path), sess_options=opts, providers=["CPUExecutionProvider"]
            )
            self._input_name = self._session.get_inputs()[0].name
        except Exception as exc:  # noqa: BLE001 - not-ready, never fatal
            self._error = f"{type(exc).__name__}: {exc}"

    @property
    def model_version(self) -> str:
        return f"bodytype:{self.version}"

    @property
    def ready(self) -> bool:
        return self._session is not None

    @property
    def error(self) -> str | None:
        return self._error

    def classify(self, frame: Any, box: BBox, plate: BBox | None) -> tuple[str, float] | None:
        """(class, probability) for the vehicle in `box`, or None when nothing could be cut."""
        if self._session is None:
            return None
        import cv2
        import numpy as np

        crop = crop_vehicle(frame, box, plate, self.crop_margin)
        if crop is None:
            return None
        resized = cv2.resize(crop, (self.input_size, self.input_size), interpolation=cv2.INTER_AREA)
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        x = np.ascontiguousarray(rgb.transpose(2, 0, 1)[None].astype(np.float32))
        logits = self._session.run(None, {self._input_name: x})[0][0]
        z = logits - logits.max()
        p = np.exp(z) / np.exp(z).sum()
        i = int(p.argmax())
        return self.classes[i], float(p[i])


class RefinedVehicleDetector:
    """Detector + classifier. The detector finds the vehicle (and picks WHICH one); when its
    class is `car` — or one the classifier was trained on — the classifier's answer replaces
    it if confident enough, else the detector's stands. A truck or bus the classifier has
    never seen is left alone: its softmax on an unknown thing means nothing."""

    def __init__(self, detector: Any, classifier: BodyTypeClassifier) -> None:
        self._detector = detector
        self._classifier = classifier
        self.stage_error: str | None = None

    @property
    def model_version(self) -> str:
        return f"{self._detector.model_version}+{self._classifier.model_version}"

    @property
    def ready(self) -> bool:
        return bool(getattr(self._detector, "ready", True))

    @property
    def error(self) -> str | None:
        parts = [
            getattr(self._detector, "error", None),
            f"classifier: {self._classifier.error}" if self._classifier.error else None,
            f"classifier: {self.stage_error}" if self.stage_error else None,
        ]
        kept = [p for p in parts if p]
        return "; ".join(kept) if kept else None

    def detect(self, image_bytes: bytes, plate: BBox | None) -> VehicleResult | None:
        import cv2
        import numpy as np

        frame = cv2.imdecode(np.frombuffer(image_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        if frame is None:
            return None
        detect_frame = getattr(self._detector, "detect_frame", None)
        base: VehicleResult | None = (
            detect_frame(frame, plate) if detect_frame else self._detector.detect(image_bytes, plate)
        )
        if base is None or base.bbox is None or not self._classifier.ready:
            return base
        if not (base.body_type == "car" or base.body_type in self._classifier.classes):
            return base
        try:
            out = self._classifier.classify(frame, base.bbox, plate)
        except Exception as exc:  # noqa: BLE001 - advisory stage, never fatal
            self.stage_error = f"{type(exc).__name__}: {exc}"
            return base
        if out is None:
            return base
        body_type, confidence = out
        if confidence < self._classifier.min_confidence:
            return base.model_copy(update={"detector_class": base.body_type})
        return base.model_copy(
            update={
                "body_type": body_type,
                "confidence": round(confidence, 4),
                "detector_class": base.body_type,
            }
        )


def time_detect(
    detector: VehicleDetector, image_bytes: bytes, plate: BBox | None
) -> tuple[VehicleResult | None, float]:
    """detect() with wall time in ms (for logs/benchmarks)."""
    started = time.perf_counter()
    result = detector.detect(image_bytes, plate)
    return result, (time.perf_counter() - started) * 1000.0