from __future__ import annotations
from pathlib import Path
import torch

class TorchTensorRTEngine:
    """TensorRT 10 engine wrapper using PyTorch CUDA pointers (zero-copy bindings)."""
    def __init__(self, path:str|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); mode=self.engine.get_tensor_mode(n)
            (self.inputs if mode==trt.TensorIOMode.INPUT else self.outputs).append(n)
    def _torch_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:dict[str,torch.Tensor]):
        for n,t in feeds.items():
            t=t.contiguous(); feeds[n]=t
            self.context.set_input_shape(n, tuple(t.shape)); self.context.set_tensor_address(n, int(t.data_ptr()))
        outs={}
        for n in self.outputs:
            shape=tuple(self.context.get_tensor_shape(n)); dt=self._torch_dtype(self.engine.get_tensor_dtype(n))
            y=torch.empty(shape,device='cuda',dtype=dt); outs[n]=y; self.context.set_tensor_address(n,int(y.data_ptr()))
        ok=self.context.execute_async_v3(stream_handle=torch.cuda.current_stream().cuda_stream)
        if not ok: raise RuntimeError('TensorRT execute_async_v3 failed')
        return outs

class DinoTRT(torch.nn.Module):
    def __init__(self, engine_path): super().__init__(); self.engine=TorchTensorRTEngine(engine_path); self.device=torch.device('cuda')
    def forward(self, pixel_values, is_training=True):
        out=self.engine.run({'pixel_values':pixel_values})['x_prenorm']; return {'x_prenorm':out}
    def to(self,*a,**k): return self
    def cpu(self): return self
    def eval(self): return self

class DsineTRT(torch.nn.Module):
    def __init__(self, engine_path): super().__init__(); self.engine=TorchTensorRTEngine(engine_path); self.device=torch.device('cuda')
    def forward(self, image, intrins):
        out=self.engine.run({'image':image,'intrins':intrins})['normal']; return [out]
    def to(self,*a,**k): return self
    def cpu(self): return self
    def eval(self): return self