from __future__ import annotations

from pathlib import Path
import math
import numpy as np
import torch


class _Allocator:
    def __new__(cls,*args,**kwargs):
        import tensorrt as trt
        class Impl(trt.IOutputAllocator):
            def __init__(self,dtype):
                trt.IOutputAllocator.__init__(self);self.dtype=dtype;self.tensor=None;self.shape=None
            def reallocate_output(self,name,memory,size,alignment):
                es=torch.empty((),dtype=self.dtype).element_size();n=max(1,(int(size)+es-1)//es)
                self.tensor=torch.empty(n,device='cuda',dtype=self.dtype);return int(self.tensor.data_ptr())
            def reallocate_output_async(self,name,memory,size,alignment,stream):return self.reallocate_output(name,memory,size,alignment)
            def notify_shape(self,name,dims):self.shape=tuple(int(x) for x in dims)
        return Impl(*args,**kwargs)


class TensorRTDDS:
    """TensorRT 11 zero-copy runner with IOutputAllocator for DDS plugin outputs."""
    def __init__(self,path):
        import tensorrt as trt
        self.trt=trt;self.logger=trt.Logger(trt.Logger.ERROR);self.runtime=trt.Runtime(self.logger);self.engine=self.runtime.deserialize_cuda_engine(Path(path).read_bytes())
        if self.engine is None:raise RuntimeError(f'Could not deserialize {path}')
        self.context=self.engine.create_execution_context()
        self.inputs=[];self.outputs=[]
        for i in range(self.engine.num_io_tensors):
            n=self.engine.get_tensor_name(i);(self.inputs if self.engine.get_tensor_mode(n)==trt.TensorIOMode.INPUT else self.outputs).append(n)
    def _dtype(self,dt):
        trt=self.trt;m={trt.float32:torch.float32,trt.float16:torch.float16,trt.int32:torch.int32,trt.int64:torch.int64,trt.bool:torch.bool}
        if hasattr(trt,'bfloat16'):m[trt.bfloat16]=torch.bfloat16
        return m[dt]
    def run(self,feeds):
        keep=[]
        for n in self.inputs:
            x=feeds[n];dt=self._dtype(self.engine.get_tensor_dtype(n));x=x.to(device='cuda',dtype=dt).contiguous();feeds[n]=x;keep.append(x);self.context.set_input_shape(n,tuple(x.shape));self.context.set_tensor_address(n,int(x.data_ptr()))
        outs={};allocs={}
        for n in self.outputs:
            dt=self._dtype(self.engine.get_tensor_dtype(n));shape=tuple(int(q) for q in self.context.get_tensor_shape(n))
            if any(q<0 for q in shape):
                a=_Allocator(dt);allocs[n]=a;self.context.set_output_allocator(n,a)
            else:
                y=torch.empty(shape if shape else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);self.context.set_tensor_address(n,int(y.data_ptr()))
        ok=self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
        if not ok:raise RuntimeError('TensorRT execute_async_v3 failed')
        torch.cuda.synchronize()
        for n,a in allocs.items():
            if a.tensor is None or a.shape is None:raise RuntimeError(f'DDS output {n} was not allocated/notified')
            num=math.prod(a.shape) if a.shape else 1;outs[n]=a.tensor[:num].view(a.shape)
        return outs