from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import tensorrt as trt


def trt_dtype_to_torch(dtype):
    if dtype == trt.DataType.HALF:
        return torch.float16
    if dtype == trt.DataType.BF16:
        return torch.bfloat16
    if dtype == trt.DataType.FLOAT:
        return torch.float32
    if dtype == trt.DataType.INT8:
        return torch.int8
    raise RuntimeError(f"Unsupported TensorRT dtype: {dtype}")


def make_inputs(engine, context, device):
    inputs = {}
    outputs = {}
    input_names = []
    output_names = []
    for idx in range(engine.num_io_tensors):
        name = engine.get_tensor_name(idx)
        shape = tuple(engine.get_tensor_shape(name))
        dtype = trt_dtype_to_torch(engine.get_tensor_dtype(name))
        if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
            context.set_input_shape(name, shape)
            input_names.append(name)
            if dtype.is_floating_point:
                tensor = torch.randn(shape, dtype=dtype, device=device)
            else:
                tensor = torch.zeros(shape, dtype=dtype, device=device)
            inputs[name] = tensor.contiguous()
        else:
            output_names.append(name)
            outputs[name] = torch.empty(shape, dtype=dtype, device=device)
    return input_names, output_names, inputs, outputs


def bind(context, names, tensors):
    for name in names:
        context.set_tensor_address(name, tensors[name].data_ptr())


def run_once(context, stream):
    ok = context.execute_async_v3(stream_handle=stream.cuda_stream)
    if not ok:
        raise RuntimeError("execute_async_v3 returned false")


def timed_loop(fn, stream, warmup, iters):
    for _ in range(warmup):
        fn()
    stream.synchronize()
    times = []
    for _ in range(iters):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record(stream)
        fn()
        end.record(stream)
        end.synchronize()
        times.append(start.elapsed_time(end))
    return times


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--engine_path", required=True)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iters", type=int, default=30)
    parser.add_argument("--cuda_graph", action="store_true")
    parser.add_argument("--json_path", type=str, default=None)
    args = parser.parse_args()

    logger = trt.Logger(trt.Logger.WARNING)
    trt.init_libnvinfer_plugins(logger, "")
    runtime = trt.Runtime(logger)
    engine_path = Path(args.engine_path)
    engine = runtime.deserialize_cuda_engine(engine_path.read_bytes())
    if engine is None:
        raise RuntimeError(f"Failed to deserialize engine: {engine_path}")
    context = engine.create_execution_context()
    device = torch.device("cuda")
    input_names, output_names, inputs, outputs = make_inputs(engine, context, device)
    bind(context, input_names, inputs)
    bind(context, output_names, outputs)
    stream = torch.cuda.Stream(device=device)

    if args.cuda_graph:
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, stream=stream):
            run_once(context, torch.cuda.current_stream(device))

        def fn():
            with torch.cuda.stream(stream):
                graph.replay()
    else:
        fn = lambda: run_once(context, stream)

    times = timed_loop(fn, stream, args.warmup, args.iters)
    result = {
        "engine_path": str(engine_path),
        "cuda_graph": args.cuda_graph,
        "warmup": args.warmup,
        "iters": args.iters,
        "mean_ms": sum(times) / len(times),
        "min_ms": min(times),
        "max_ms": max(times),
        "times_ms": times,
        "inputs": {name: {"shape": list(inputs[name].shape), "dtype": str(inputs[name].dtype)} for name in input_names},
        "outputs": {name: {"shape": list(outputs[name].shape), "dtype": str(outputs[name].dtype)} for name in output_names},
    }
    print(json.dumps(result, indent=2))
    if args.json_path:
        out = Path(args.json_path)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()