#!/usr/bin/env python3
"""
Edge inference server — Ascend NPU (CANN ACL) backend.
Replaces ultralytics/PyTorch with ACL for Atlas 200I DK2.
Socket protocol unchanged: 4-byte big-endian length prefix + JSON body.

Output formats (OUTPUT_FORMAT env):
  raw      — YOLOv8 style [1, 4+nc, anchors], NMS applied in Python
  nms_free — YOLOv10 style [1, topk, 6], already NMS'd by model
"""
import os
import json
import struct
import socket
import base64
import logging
import numpy as np
import cv2
import acl

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

SOCK_PATH   = os.getenv("INFER_SOCKET",    "/tmp/edge-infer.sock")
MODEL_PATH  = os.getenv("OM_MODEL",        "model.om")
CONF_TH     = float(os.getenv("CONF_THRESHOLD", "0.5"))
IOU_TH      = float(os.getenv("IOU_THRESHOLD",  "0.45"))
DEVICE_ID   = int(os.getenv("DEVICE_ID",   "0"))
NAMES_FILE  = os.getenv("NAMES_FILE",      "")
OUTPUT_FMT  = os.getenv("OUTPUT_FORMAT",   "raw")   # raw | nms_free

ACL_MEM_MALLOC_NORMAL_ONLY = 0
ACL_MEMCPY_HOST_TO_DEVICE  = 1
ACL_MEMCPY_DEVICE_TO_HOST  = 2


def load_names(path):
    if path and os.path.exists(path):
        with open(path) as f:
            return [l.strip() for l in f if l.strip()]
    return [str(i) for i in range(1000)]


class AclModel:
    def __init__(self, model_path, device_id):
        self.device_id = device_id
        self._init_acl()
        self._load_model(model_path)
        self._alloc_outputs()
        log.info("model loaded  path=%s  input=%s  outputs=%d",
                 model_path, self.input_shape, self.output_num)

    def _init_acl(self):
        ret = acl.init()
        assert ret == 0, f"acl.init failed ret={ret}"
        ret = acl.rt.set_device(self.device_id)
        assert ret == 0, f"set_device failed ret={ret}"
        self.context, ret = acl.rt.create_context(self.device_id)
        assert ret == 0, f"create_context failed ret={ret}"

    def _load_model(self, path):
        self.model_id, ret = acl.mdl.load_from_file(path)
        assert ret == 0, f"load_from_file failed ret={ret}"

        self.desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.desc, self.model_id)
        assert ret == 0

        self.input_num  = acl.mdl.get_num_inputs(self.desc)
        self.output_num = acl.mdl.get_num_outputs(self.desc)

        dims, ret = acl.mdl.get_input_dims(self.desc, 0)
        assert ret == 0
        self.input_shape = list(dims["dims"])   # [1, 3, H, W]
        self.input_h     = self.input_shape[2]
        self.input_w     = self.input_shape[3]
        self.input_size  = acl.mdl.get_input_size_by_index(self.desc, 0)

        self.output_shapes = []
        for i in range(self.output_num):
            d, ret = acl.mdl.get_output_dims(self.desc, i)
            assert ret == 0
            self.output_shapes.append(list(d["dims"]))

    def _alloc_outputs(self):
        self.out_bufs  = []
        self.out_sizes = []
        for i in range(self.output_num):
            sz = acl.mdl.get_output_size_by_index(self.desc, i)
            buf, ret = acl.rt.malloc(sz, ACL_MEM_MALLOC_NORMAL_ONLY)
            assert ret == 0
            self.out_bufs.append(buf)
            self.out_sizes.append(sz)

    def preprocess(self, img_bgr):
        """Letterbox → RGB → NCHW float32 [0,1]. Returns blob, scale, pad_top, pad_left."""
        h0, w0 = img_bgr.shape[:2]
        scale = min(self.input_h / h0, self.input_w / w0)
        nh, nw = int(h0 * scale), int(w0 * scale)
        resized = cv2.resize(img_bgr, (nw, nh), interpolation=cv2.INTER_LINEAR)

        canvas = np.full((self.input_h, self.input_w, 3), 114, dtype=np.uint8)
        pad_top  = (self.input_h - nh) // 2
        pad_left = (self.input_w - nw) // 2
        canvas[pad_top:pad_top + nh, pad_left:pad_left + nw] = resized

        rgb  = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
        blob = rgb.astype(np.float32) / 255.0
        blob = np.ascontiguousarray(blob.transpose(2, 0, 1)[np.newaxis])  # NCHW
        return blob, scale, pad_top, pad_left

    def run(self, blob):
        """Push blob to NPU, execute, pull outputs back as numpy arrays."""
        in_ds = acl.mdl.create_dataset()
        in_buf, ret = acl.rt.malloc(self.input_size, ACL_MEM_MALLOC_NORMAL_ONLY)
        assert ret == 0
        ret = acl.rt.memcpy(in_buf, self.input_size,
                            blob.tobytes(), self.input_size,
                            ACL_MEMCPY_HOST_TO_DEVICE)
        assert ret == 0
        db = acl.create_data_buffer(in_buf, self.input_size)
        _, ret = acl.mdl.add_dataset_buffer(in_ds, db)
        assert ret == 0

        out_ds = acl.mdl.create_dataset()
        for i in range(self.output_num):
            db = acl.create_data_buffer(self.out_bufs[i], self.out_sizes[i])
            _, ret = acl.mdl.add_dataset_buffer(out_ds, db)
            assert ret == 0

        ret = acl.mdl.execute(self.model_id, in_ds, out_ds)
        assert ret == 0, f"mdl.execute failed ret={ret}"

        outputs = []
        for i in range(self.output_num):
            sz   = self.out_sizes[i]
            host = np.zeros(sz, dtype=np.uint8)
            ret  = acl.rt.memcpy(host.ctypes.data, sz,
                                 self.out_bufs[i], sz,
                                 ACL_MEMCPY_DEVICE_TO_HOST)
            assert ret == 0
            outputs.append(host.view(np.float32).reshape(self.output_shapes[i]))

        acl.rt.free(in_buf)
        acl.mdl.destroy_dataset(in_ds)
        acl.mdl.destroy_dataset(out_ds)
        return outputs

    def destroy(self):
        for buf in self.out_bufs:
            acl.rt.free(buf)
        acl.mdl.unload(self.model_id)
        acl.mdl.destroy_desc(self.desc)
        acl.rt.destroy_context(self.context)
        acl.rt.reset_device(self.device_id)
        acl.finalize()


# ------------------------------------------------------------------ #
# Post-processing
# ------------------------------------------------------------------ #

def _xywh2xyxy(boxes):
    out = np.empty_like(boxes)
    out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
    out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
    out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
    out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
    return out


def _unpad(x1, y1, x2, y2, scale, pad_top, pad_left, orig_h, orig_w):
    x1 = max(0.0, (x1 - pad_left) / scale)
    y1 = max(0.0, (y1 - pad_top)  / scale)
    x2 = min(float(orig_w), (x2 - pad_left) / scale)
    y2 = min(float(orig_h), (y2 - pad_top)  / scale)
    return x1, y1, x2, y2


def postprocess_raw(output, conf_th, iou_th, scale, pad_top, pad_left, orig_h, orig_w):
    """YOLOv8 raw output [1, 4+nc, anchors] → detections list."""
    pred    = output[0].T          # [anchors, 4+nc]
    boxes   = pred[:, :4]          # cx,cy,w,h in input coords
    scores  = pred[:, 4:]

    cls_ids = scores.argmax(axis=1)
    confs   = scores[np.arange(len(scores)), cls_ids]

    mask = confs >= conf_th
    boxes, confs, cls_ids = boxes[mask], confs[mask], cls_ids[mask]
    if len(boxes) == 0:
        return []

    xyxy = _xywh2xyxy(boxes)
    keep = cv2.dnn.NMSBoxes(xyxy.tolist(), confs.tolist(), conf_th, iou_th)
    if len(keep) == 0:
        return []

    dets = []
    for idx in np.array(keep).flatten():
        x1, y1, x2, y2 = _unpad(*xyxy[idx], scale, pad_top, pad_left, orig_h, orig_w)
        dets.append({"class_id": int(cls_ids[idx]),
                     "conf": float(confs[idx]),
                     "bbox": [x1, y1, x2, y2]})
    return dets


def postprocess_nms_free(output, conf_th, scale, pad_top, pad_left, orig_h, orig_w):
    """YOLOv10 NMS-free output [1, topk, 6] (x1,y1,x2,y2,conf,cls) → detections list."""
    dets = []
    for row in output[0]:
        x1, y1, x2, y2, conf, cls_id = row
        if conf < conf_th:
            continue
        x1, y1, x2, y2 = _unpad(x1, y1, x2, y2, scale, pad_top, pad_left, orig_h, orig_w)
        dets.append({"class_id": int(cls_id),
                     "conf": float(conf),
                     "bbox": [x1, y1, x2, y2]})
    return dets


# ------------------------------------------------------------------ #
# Socket helpers
# ------------------------------------------------------------------ #

def recv_msg(conn):
    hdr = conn.recv(4)
    if not hdr:
        return None
    length = struct.unpack(">I", hdr)[0]
    data = b""
    while len(data) < length:
        chunk = conn.recv(length - len(data))
        if not chunk:
            return None
        data += chunk
    return data


def send_msg(conn, payload):
    data = json.dumps(payload).encode("utf-8")
    conn.sendall(struct.pack(">I", len(data)) + data)


# ------------------------------------------------------------------ #
# Inference entry
# ------------------------------------------------------------------ #

def infer_one(model, names, msg):
    jpg = base64.b64decode(msg["jpeg_b64"])
    arr = np.frombuffer(jpg, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    orig_h, orig_w = img.shape[:2]

    blob, scale, pad_top, pad_left = model.preprocess(img)
    outputs = model.run(blob)

    if OUTPUT_FMT == "nms_free":
        raw_dets = postprocess_nms_free(
            outputs[0], CONF_TH, scale, pad_top, pad_left, orig_h, orig_w)
    else:
        raw_dets = postprocess_raw(
            outputs[0], CONF_TH, IOU_TH, scale, pad_top, pad_left, orig_h, orig_w)

    dets = [{"class": names[d["class_id"]] if d["class_id"] < len(names) else str(d["class_id"]),
             "conf":  d["conf"],
             "bbox":  d["bbox"]} for d in raw_dets]

    return {"stream_id": msg["stream_id"],
            "device_id": msg["device_id"],
            "ts":        msg["ts"],
            "detections": dets}


def main():
    names = load_names(NAMES_FILE)
    model = AclModel(MODEL_PATH, DEVICE_ID)

    if os.path.exists(SOCK_PATH):
        os.remove(SOCK_PATH)
    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(SOCK_PATH)
    srv.listen(4)
    log.info("infer server ready  socket=%s  device=%d  fmt=%s",
             SOCK_PATH, DEVICE_ID, OUTPUT_FMT)

    try:
        while True:
            conn, _ = srv.accept()
            try:
                while True:
                    data = recv_msg(conn)
                    if data is None:
                        break
                    msg = json.loads(data.decode("utf-8"))
                    out = infer_one(model, names, msg)
                    send_msg(conn, out)
            finally:
                conn.close()
    finally:
        model.destroy()


if __name__ == "__main__":
    main()