"""Shared ONNX Runtime inference for the PCB defect detector.

Used by both the FastAPI service (serve/) and the Gradio demo (app/app.py), so the
two always run the same preprocessing, decoding and NMS. The steps mirror
Ultralytics 8.3 (LetterBox, non_max_suppression, scale_boxes) so results match the
evaluation pipeline in src/.

Only numpy, OpenCV and onnxruntime are needed: no PyTorch, no Ultralytics.
"""
from __future__ import annotations

import ast
import time
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort

FALLBACK_NAMES = ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"]
COLORS = [(56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255), (49, 210, 207), (10, 249, 72)]  # BGR
MAX_WH = 7680  # class offset for batched NMS, as in Ultralytics


@dataclass(frozen=True)
class Detection:
    class_id: int
    class_name: str
    confidence: float
    box: tuple[float, float, float, float]  # x1, y1, x2, y2 in original-image pixels

    def to_dict(self) -> dict:
        return {
            "class": self.class_name,
            "class_id": self.class_id,
            "confidence": round(self.confidence, 4),
            "box": [round(v, 1) for v in self.box],
        }


@dataclass(frozen=True)
class Timings:
    preprocess_ms: float
    model_ms: float
    postprocess_ms: float

    @property
    def total_ms(self) -> float:
        return self.preprocess_ms + self.model_ms + self.postprocess_ms


def letterbox(bgr: np.ndarray, size: int) -> tuple[np.ndarray, float, tuple[float, float]]:
    """Ultralytics LetterBox(auto=False, center=True): resize keeping aspect ratio, pad with 114."""
    h, w = bgr.shape[:2]
    r = min(size / h, size / w)
    new_w, new_h = int(round(w * r)), int(round(h * r))
    dw, dh = (size - new_w) / 2, (size - new_h) / 2
    img = bgr if (w, h) == (new_w, new_h) else cv2.resize(bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
    return img, r, (dw, dh)


def nms(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> np.ndarray:
    """Greedy NMS (same rule as torchvision.ops.nms). boxes: (n, 4) xyxy. Returns kept indices."""
    order = scores.argsort()[::-1]
    x1, y1, x2, y2 = boxes.T
    areas = (x2 - x1) * (y2 - y1)
    keep = []
    while order.size:
        i = order[0]
        keep.append(i)
        rest = order[1:]
        iw = np.clip(np.minimum(x2[i], x2[rest]) - np.maximum(x1[i], x1[rest]), 0, None)
        ih = np.clip(np.minimum(y2[i], y2[rest]) - np.maximum(y1[i], y1[rest]), 0, None)
        inter = iw * ih
        iou = inter / (areas[i] + areas[rest] - inter + 1e-9)
        order = rest[iou <= iou_thres]
    return np.asarray(keep, dtype=int)


class Detector:
    """Loads the ONNX model once; predict() is safe to call from several threads."""

    def __init__(self, model_path: str | Path, threads: int | None = None):
        self.model_path = Path(model_path)
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        if threads:
            opts.intra_op_num_threads = threads
            opts.inter_op_num_threads = 1
        self.session = ort.InferenceSession(str(self.model_path), opts, providers=["CPUExecutionProvider"])
        inp = self.session.get_inputs()[0]
        self.input_name = inp.name
        meta = self.session.get_modelmeta().custom_metadata_map
        self.imgsz = inp.shape[2] if isinstance(inp.shape[2], int) else ast.literal_eval(meta.get("imgsz", "[1280]"))[0]
        names = ast.literal_eval(meta["names"]) if "names" in meta else dict(enumerate(FALLBACK_NAMES))
        self.names = {int(k): v for k, v in names.items()}

    # --- pipeline steps -------------------------------------------------------
    def preprocess(self, bgr: np.ndarray):
        img, r, pad = letterbox(bgr, self.imgsz)
        x = np.ascontiguousarray(img[:, :, ::-1].transpose(2, 0, 1), dtype=np.float32)[None] / 255.0
        return x, r, pad

    def infer(self, x: np.ndarray) -> np.ndarray:
        return self.session.run(None, {self.input_name: x})[0]

    def postprocess(self, raw, r, pad, shape, conf=0.25, iou=0.7, max_det=300) -> list[Detection]:
        pred = raw[0].T                                    # (anchors, 4 + num_classes)
        scores_all = pred[:, 4:]
        cls = scores_all.argmax(1)
        scores = scores_all[np.arange(len(cls)), cls]
        keep = scores > conf
        if not keep.any():
            return []
        xywh, cls, scores = pred[keep, :4], cls[keep], scores[keep]
        boxes = np.empty_like(xywh)
        boxes[:, 0] = xywh[:, 0] - xywh[:, 2] / 2
        boxes[:, 1] = xywh[:, 1] - xywh[:, 3] / 2
        boxes[:, 2] = xywh[:, 0] + xywh[:, 2] / 2
        boxes[:, 3] = xywh[:, 1] + xywh[:, 3] / 2
        idx = nms(boxes + cls[:, None] * MAX_WH, scores, iou)[:max_det]
        boxes, cls, scores = boxes[idx], cls[idx], scores[idx]

        # Ultralytics scale_boxes: undo padding and resize, then clip to the image.
        h, w = shape
        pad_x = round((self.imgsz - w * r) / 2 - 0.1)
        pad_y = round((self.imgsz - h * r) / 2 - 0.1)
        boxes = (boxes - [pad_x, pad_y, pad_x, pad_y]) / r
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
        return [
            Detection(int(c), self.names.get(int(c), str(c)), float(s), tuple(float(v) for v in b))
            for b, c, s in zip(boxes, cls, scores)
        ]

    def predict(self, bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7) -> tuple[list[Detection], Timings]:
        t0 = time.perf_counter()
        x, r, pad = self.preprocess(bgr)
        t1 = time.perf_counter()
        raw = self.infer(x)
        t2 = time.perf_counter()
        dets = self.postprocess(raw, r, pad, bgr.shape[:2], conf, iou)
        t3 = time.perf_counter()
        return dets, Timings((t1 - t0) * 1000, (t2 - t1) * 1000, (t3 - t2) * 1000)


def draw(bgr: np.ndarray, dets: list[Detection]) -> np.ndarray:
    """Return a copy of the image with boxes and labels drawn on it."""
    out = bgr.copy()
    thick = max(2, round(max(out.shape[:2]) / 600))
    for d in dets:
        x1, y1, x2, y2 = (int(v) for v in d.box)
        color = COLORS[d.class_id % len(COLORS)]
        cv2.rectangle(out, (x1, y1), (x2, y2), color, thick)
        cv2.putText(out, f"{d.class_name} {d.confidence:.2f}", (x1, max(20, y1 - 8)),
                    cv2.FONT_HERSHEY_SIMPLEX, thick / 2.5, color, thick, cv2.LINE_AA)
    return out