"""TensorRT Acceleration Utilities for RIFE.

This module provides the low-level interface for NVIDIA TensorRT, handling
engine compilation, timing caches, memory allocation, and optimized inference 
using CUDA Graphs.
"""

import copy
import io
import os

import cuda.bindings.runtime as cudart
import onnx
import numpy as np
import tensorrt_rtx as trt
import torch

from collections import OrderedDict
from logging import error, warning
from onnxconverter_common import float16
from onnxsim import simplify

# Use polygraphy TensorRTX backend
os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1"

from polygraphy.backend.common import bytes_from_path
from polygraphy import util
from polygraphy.backend.trt import ModifyNetworkOutputs, Profile
from polygraphy.backend.trt import (
    engine_from_bytes,
    engine_from_network,
    network_from_onnx_path,
    network_from_onnx_bytes,
    save_engine,
)
from polygraphy.logger import G_LOGGER

from torch.cuda import nvtx
from tqdm import tqdm

G_LOGGER.module_severity = G_LOGGER.ERROR

# Map of numpy dtype -> torch dtype
numpy_to_torch_dtype_dict = {
    np.uint8: torch.uint8,
    np.int8: torch.int8,
    np.int16: torch.int16,
    np.int32: torch.int32,
    np.int64: torch.int64,
    np.float16: torch.float16,
    np.float32: torch.float32,
    np.float64: torch.float64,
    np.complex64: torch.complex64,
    np.complex128: torch.complex128,
}
if np.version.full_version >= "1.24.0":
    numpy_to_torch_dtype_dict[np.bool_] = torch.bool
else:
    numpy_to_torch_dtype_dict[np.bool] = torch.bool

# Map of torch dtype -> numpy dtype
torch_to_numpy_dtype_dict = {
    value: key for (key, value) in numpy_to_torch_dtype_dict.items()
}

# Path for your persistent hardware profile
node_dir = os.path.dirname(os.path.realpath(__file__))
MASTER_CACHE_FILE = f"{node_dir}\\RIFE_RTX_ENGINE.cache"

def get_timing_cache(config):
    """Retrieves or creates a TensorRT timing cache.

    Args:
        config (trt.IBuilderConfig): The builder configuration object.

    Returns:
        trt.ITimingCache: The loaded or new timing cache.
    """
    if os.path.exists(MASTER_CACHE_FILE):
        with open(MASTER_CACHE_FILE, "rb") as f:
            cache_data = f.read()
            return config.create_timing_cache(cache_data)
    return config.create_timing_cache(b"")

def save_timing_cache(timing_cache):
    """Saves the current timing cache to disk for future builds.

    Args:
        timing_cache (trt.ITimingCache): The cache to serialize.
    """
    with open(MASTER_CACHE_FILE, "wb") as f:
        f.write(timing_cache.serialize())
 
def CUASSERT(cuda_ret):
    """Asserts that a CUDA call was successful.
    https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py
    
    Args:
        cuda_ret (tuple): The return value from a cuda-python call.

    Returns:
        Any: The result of the CUDA call if successful.

    Raises:
        RuntimeError: If the CUDA call returned an error.
    """
    err = cuda_ret[0]
    if err != cudart.cudaError_t.cudaSuccess:
        raise RuntimeError(
            f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
        )
    if len(cuda_ret) > 1:
        return cuda_ret[1]
    return None

class TQDMProgressMonitor(trt.IProgressMonitor):
    """Progress monitor for TensorRT engine builds using TQDM bars."""
    
    def __init__(self):
        trt.IProgressMonitor.__init__(self)
        self._active_phases = {}
        self._step_result = True
        self.max_indent = 5

    def phase_start(self, phase_name, parent_phase, num_steps):
        """Called when a build phase starts."""
        leave = False
        try:
            if parent_phase is not None:
                nbIndents = (
                    self._active_phases.get(parent_phase, {}).get(
                        "nbIndents", self.max_indent
                    )
                    + 1
                )
                if nbIndents >= self.max_indent:
                    return
            else:
                nbIndents = 0
                leave = True
            self._active_phases[phase_name] = {
                "tq": tqdm(
                    total=num_steps, desc=phase_name, leave=leave, position=nbIndents
                ),
                "nbIndents": nbIndents,
                "parent_phase": parent_phase,
            }
        except KeyboardInterrupt:
            # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
            _step_result = False

    def phase_finish(self, phase_name):
        """Called when a build phase finishes."""
        try:
            if phase_name in self._active_phases.keys():
                self._active_phases[phase_name]["tq"].update(
                    self._active_phases[phase_name]["tq"].total
                    - self._active_phases[phase_name]["tq"].n
                )

                parent_phase = self._active_phases[phase_name].get("parent_phase", None)
                while parent_phase is not None:
                    self._active_phases[parent_phase]["tq"].refresh()
                    parent_phase = self._active_phases[parent_phase].get(
                        "parent_phase", None
                    )
                if (
                    self._active_phases[phase_name]["parent_phase"]
                    in self._active_phases.keys()
                ):
                    self._active_phases[
                        self._active_phases[phase_name]["parent_phase"]
                    ]["tq"].refresh()
                del self._active_phases[phase_name]
            pass
        except KeyboardInterrupt:
            _step_result = False

    def step_complete(self, phase_name, step):
        """Called when a step within a phase completes."""
        try:
            if phase_name in self._active_phases.keys():
                self._active_phases[phase_name]["tq"].update(
                    step - self._active_phases[phase_name]["tq"].n
                )
            return self._step_result
        except KeyboardInterrupt:
            # There is no need to propagate this exception to TensorRT. We can simply cancel the build.
            return False

class Engine:
    """Wrapper for TensorRT engine management.

    Handles loading, building, and executing inference with optional 
    CUDA Graph support.

    Attributes:
        engine_path (str): File path to the .trt engine.
    """
    
    def __init__(
        self,
        engine_path,
    ):
        self.engine_path = engine_path
        self.engine = None
        self.context = None
        self.buffers = OrderedDict()
        self.tensors = OrderedDict()
        self.cuda_graph_instance = None  # cuda graph
        self.graph = None

    def __del__(self):
        """Ensures CUDA resources are freed on deletion."""
        if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
            try:
                cudart.cudaGraphDestroy(self.cuda_graph_instance)
            except:
                pass
        if hasattr(self, 'graph') and self.graph is not None:
            try:
                cudart.cudaGraphDestroy(self.graph)
            except:
                pass

        del self.engine
        del self.context
        del self.buffers
        del self.tensors

    def reset(self, engine_path=None):
        """Cleans up engine and CUDA resources.

        Args:
            engine_path (str, optional): New path to set for the engine.
        """
        if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
            try:
                cudart.cudaGraphDestroy(self.cuda_graph_instance)
            except:
                pass
            self.cuda_graph_instance = None
        if hasattr(self, 'graph') and self.graph is not None:
            try:
                cudart.cudaGraphDestroy(self.graph)
            except:
                pass
            self.graph = None

        if hasattr(self, 'engine') and self.engine is not None:
            del self.engine
        if hasattr(self, 'context') and self.context is not None:
            del self.context
        if hasattr(self, 'buffers'):
            del self.buffers
        if hasattr(self, 'tensors'):
            del self.tensors

        self.engine = None
        self.context = None
        self.engine_path = engine_path if engine_path else self.engine_path

        self.buffers = OrderedDict()
        self.tensors = OrderedDict()
        self.inputs = {}
        self.outputs = {}

    def build(
        self,
        onnx_path,
        input_profile=None,
        enable_refit=False,
        enable_preview=False,
        enable_all_tactics=False,
        timing_cache=None,
        update_output_names=None,
    ):
        """Compiles an ONNX model into a TensorRT engine.

        Args:
            onnx_path (str): Path to source ONNX file.
            input_profile (list): List of dicts defining min/opt/max shapes.
            **kwargs: Configuration flags (enable_refit, etc).

        Returns:
            int: 0 if success, 1 if failure.
        """
        p = [Profile()]
        if input_profile:
            p = [Profile() for i in range(len(input_profile))]
            for _p, i_profile in zip(p, input_profile):
                for name, dims in i_profile.items():
                    assert len(dims) == 3
                    _p.add(name, min=dims[0], opt=dims[1], max=dims[2])

        config_kwargs = {}
        if not enable_all_tactics:
            config_kwargs["tactic_sources"] = []  

        # Load ONNX file then convert the model to fp16
        # Probably might cause issues if onnx file given is quantized or whatever
        onnx_bytes       = onnx.load(onnx_path)
        onnx_precision   = onnx_bytes.graph.input[0].type.tensor_type.elem_type        
        onnx_simp, check = simplify(float16.convert_float_to_float16(onnx_bytes) if onnx_precision == onnx.TensorProto.FLOAT else onnx_bytes)
        
        if not check:
            error(f"Simplified ONNX model: {onnx_path} could not be validated")
            return 1
        
        # Keep simplified model in memory to not overwrite the onnx file in models directory
        # We dont want to be destructive
        onnx_buffer = io.BytesIO()
        onnx_buffer.write(onnx_simp.SerializeToString())
        onnx_buffer.seek(0)
        onnx_buffer.name = "onnx_simp_mem_rife.onnx" 
        
        #network = network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM])
        network = network_from_onnx_bytes(onnx_buffer.getvalue(), flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM])
        
        if update_output_names:
            print(f"Updating network outputs to {update_output_names}")
            network = ModifyNetworkOutputs(network, update_output_names)

        builder = network[0]
        config = builder.create_builder_config()
        config.progress_monitor = TQDMProgressMonitor()

        #config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
        config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None
        
        # Limit workspace (e.g., to 2GB or 4GB depending on your GPU)
        # 1 << 30 is 1GB
        #config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1 << 30))
        config.builder_optimization_level = 5
        
        # Use Master Cache
        timing_cache = get_timing_cache(config)
        config.set_timing_cache(timing_cache, False) 

        profiles = copy.deepcopy(p)
        for profile in profiles:
            # Last profile is used for set_calibration_profile.
            calib_profile = profile.fill_defaults(network[1]).to_trt(
                builder, network[1]
            )
            config.add_optimization_profile(calib_profile)

        result = 0

        try:
            engine = engine_from_network(network, config)
            save_engine(engine, path=self.engine_path)
        except Exception as e:
            error(f"Failed to build engine: {e}")            
            result = 1
            
        # Update Master Cache with new findings
        save_timing_cache(timing_cache)
        print(f"Updated Master Cache.")
        
        return result

    def load(self):
        """Loads the serialized engine from disk."""
        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))

    def activate(self, reuse_device_memory=None):
        """Initializes the execution context."""
        if self.engine is None:
            self.load()

        if reuse_device_memory:
            self.context = self.engine.create_execution_context_without_device_memory()
        #    self.context.device_memory = reuse_device_memory
        else:
            self.context = self.engine.create_execution_context()

    def allocate_buffers(self, shape_dict=None, device="cuda"):
        """Allocates GPU memory for input and output tensors.

        Args:
            shape_dict (dict, optional): Specific shapes for tensors.
            device (str): Device to allocate on (default "cuda").
        """
        if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
            try:
                cudart.cudaGraphDestroy(self.cuda_graph_instance)
            except:
                pass
            self.cuda_graph_instance = None
        if hasattr(self, 'graph') and self.graph is not None:
            try:
                cudart.cudaGraphDestroy(self.graph)
            except:
                pass
            self.graph = None

        nvtx.range_push("allocate_buffers")
        for idx in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(idx)
            binding = self.engine[idx]
            if shape_dict and binding in shape_dict:
                shape = shape_dict[binding]["shape"]
            else:
                shape = self.context.get_tensor_shape(name)

            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.context.set_input_shape(name, shape)
            tensor = torch.empty(
                tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
            ).to(device=device)
            self.tensors[binding] = tensor
        nvtx.range_pop()

    def infer(self, feed_dict, stream, use_cuda_graph=False):
        """Performs inference on the provided data.

        Args:
            feed_dict (dict): Input tensors.
            stream (cuda.Stream): CUDA stream for execution.
            use_cuda_graph (bool): If True, uses/captures CUDA graphs.

        Returns:
            OrderedDict: The output tensors.
        """
        for name, buf in feed_dict.items():
            self.tensors[name].copy_(buf)

        for name, tensor in self.tensors.items():
            self.context.set_tensor_address(name, tensor.data_ptr())

        if use_cuda_graph:
            if self.cuda_graph_instance is not None:
                CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr))
                CUASSERT(cudart.cudaStreamSynchronize(stream.ptr))
            else:
                # do inference before CUDA graph capture
                noerror = self.context.execute_async_v3(stream.ptr)
                if not noerror:
                    raise ValueError("ERROR: inference failed.")
                # capture cuda graph
                CUASSERT(
                    cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
                )
                self.context.execute_async_v3(stream.ptr)
                self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr))
                self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
        else:
            noerror = self.context.execute_async_v3(stream.ptr)
            if not noerror:
                raise ValueError("ERROR: inference failed.")

        return self.tensors