"""استدلال محلي على الجهاز لنماذج edge_model (ONNX Runtime — خفيف، بلا AGPL).

يشغّل نموذج YOLOv8-ONNX على لقطة الكاميرا ويُرجع الأجسام المكتشفة.
نماذج القطاعات (PPE، أسلحة، جوال…) تدخل نفس الحلقة بمجرد تصديرها إلى ONNX
+ قائمة أصنافها؛ لا تغيير كود. النموذج الافتراضي COCO يكشف person/cell phone/knife…
"""

import numpy as np

COCO80 = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
    "hair drier", "toothbrush",
]


class OnnxDetector:
    """يحمّل نموذج ONNX مرة ويشغّله على صور JPEG. تحميل كسول للاعتماديات الثقيلة."""

    def __init__(self, model_path: str, labels=None, imgsz: int = 640):
        import onnxruntime as ort
        self.sess = ort.InferenceSession(
            model_path, providers=["CPUExecutionProvider"])
        self.inp = self.sess.get_inputs()[0].name
        # أصناف النموذج: تُقرأ ذاتياً من بيانات ONNX الوصفية (Ultralytics تضمّنها)،
        # فأي نموذج قطاعي مصدَّر يوسم صحيحاً بلا ملف أصناف خارجي — وإلا COCO افتراضياً.
        self.labels = labels or self._labels_from_meta() or COCO80
        self.imgsz = imgsz

    def _labels_from_meta(self):
        import ast
        meta = self.sess.get_modelmeta().custom_metadata_map or {}
        raw = meta.get("names")
        if not raw:
            return None
        try:
            names = ast.literal_eval(raw)          # "{0: 'person', …}" → dict
            return [names[i] for i in range(len(names))]
        except Exception:
            return None

    def _letterbox(self, img):
        from PIL import Image
        w0, h0 = img.size
        r = min(self.imgsz / w0, self.imgsz / h0)
        nw, nh = round(w0 * r), round(h0 * r)
        img = img.resize((nw, nh), Image.BILINEAR)
        canvas = Image.new("RGB", (self.imgsz, self.imgsz), (114, 114, 114))
        px, py = (self.imgsz - nw) // 2, (self.imgsz - nh) // 2
        canvas.paste(img, (px, py))
        return canvas, r, px, py

    def detect(self, jpeg_bytes: bytes, conf=0.35, iou=0.45):
        """يُرجع [(label, confidence, (x1,y1,x2,y2)), …] بإحداثيات الصورة الأصلية."""
        import io
        from PIL import Image
        img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
        w0, h0 = img.size
        lb, r, px, py = self._letterbox(img)
        x = np.asarray(lb, dtype=np.float32).transpose(2, 0, 1)[None] / 255.0
        out = self.sess.run(None, {self.inp: x})[0]      # [1,84,8400]
        pred = np.squeeze(out).T                          # [8400,84]
        cls = pred[:, 4:]
        conf_scores = cls.max(1)
        keep = conf_scores > conf
        pred, conf_scores, cls = pred[keep], conf_scores[keep], cls[keep]
        if not len(pred):
            return []
        cls_ids = cls.argmax(1)
        cx, cy, bw, bh = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]
        boxes = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], 1)
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - px) / r    # فكّ letterbox → أصلي
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - py) / r
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w0)
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h0)
        idx = _nms(boxes, conf_scores, iou)
        res = []
        for i in idx:
            name = self.labels[cls_ids[i]] if cls_ids[i] < len(self.labels) else str(cls_ids[i])
            res.append((name, float(conf_scores[i]), tuple(boxes[i].round().astype(int))))
        return res


def _nms(boxes, scores, iou_thr):
    x1, y1, x2, y2 = boxes.T
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        ovr = inter / (areas[i] + areas[order[1:]] - inter + 1e-9)
        order = order[1:][ovr <= iou_thr]
    return keep


if __name__ == "__main__":  # self-check: يكشف person في صورة اختبار
    import sys
    det = OnnxDetector(sys.argv[1])
    with open(sys.argv[2], "rb") as f:
        results = det.detect(f.read())
    labels = {r[0] for r in results}
    print("detected:", sorted(labels))
    assert results, "expected at least one detection"
    print(f"OK — {len(results)} objects")