#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import ctypes
from typing import Optional, List, Union

import numpy as np
import tensorrt as trt

from cuda.bindings import driver as cuda, runtime as cudart, nvrtc


class ArrayWithOwner(np.ndarray):
    """Numpy array that holds a reference to its owner object"""
    def __new__(cls, input_array, owner):
        obj = np.asarray(input_array).view(cls)
        obj._owner = owner
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        self._owner = getattr(obj, '_owner', None)



def cuda_call(call):
    """Helper function to make CUDA calls and check for errors"""
    def _cudaGetErrorEnum(error):
        if isinstance(error, cuda.CUresult):
            err, name = cuda.cuGetErrorName(error)
            return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
        elif isinstance(error, cudart.cudaError_t):
            return cudart.cudaGetErrorName(error)[1]
        elif isinstance(error, nvrtc.nvrtcResult):
            return nvrtc.nvrtcGetErrorString(error)[1]
        else:
            raise RuntimeError("Unknown error type: {}".format(error))

    err, res = call[0], call[1:]
    if err.value:
        raise RuntimeError(
            "CUDA error code={}({})".format(
                err.value, _cudaGetErrorEnum(err)
            )
        )
    if len(res) == 1:
        return res[0]
    elif len(res) == 0:
        return None
    else:
        return res


def create_cuda_context(device):
    """
    Create CUDA context with version-aware API handling.

    Handles different CUDA API versions based on actual documented signatures:
    - CUDA 11.8-12.9: cuCtxCreate(flags, device) - 2 arguments
    - CUDA 13.0+: cuCtxCreate(ctxCreateParams, flags, device) - 3 arguments

    Args:
        device: CUDA device handle from cuDeviceGet

    Returns:
        CUDA context handle
    """
    # Try different API versions
    try:
        # Try CUDA 13.0+ API first (3 arguments with ctxCreateParams)
        # cuCtxCreate(ctxCreateParams, flags, device)
        return cuda_call(cuda.cuCtxCreate(None, 0, device))
    except TypeError:
        # CUDA 11.8-12.9 API: cuCtxCreate(flags, device)
        return cuda_call(cuda.cuCtxCreate(0, device))



class HostDeviceMem:
    """Pair of host and device memory using RAII composition"""
    def __init__(self, size: int, dtype: Optional[np.dtype] = None):
        if dtype is None:
            dtype = np.dtype(np.uint8)
        else:
            dtype = np.dtype(dtype)
        self._size = size
        self._dtype = dtype

        # Use RAII classes for memory management
        self._host_mem = PinnedHostMem(size, dtype)
        self._device_mem = DeviceMem(size * dtype.itemsize)

    @property
    def host(self) -> np.ndarray:
        # Return the array directly - ArrayWithOwner ensures proper lifetime management
        return self._host_mem.array

    @host.setter
    def host(self, data: Union[np.ndarray, bytes]):
        # Delegate to PinnedHostMem for proper data handling
        self._host_mem.array = data

    @property
    def device_ptr(self) -> int:
        """Device memory pointer"""
        return self._device_mem.device_ptr

    @property
    def nbytes(self) -> int:
        return self._host_mem.nbytes


    def __str__(self):
        return f"Host:\n{self.host}\nDevice:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"

    def __repr__(self):
        return self.__str__()


class DeviceMem:
    """Device-only memory allocation for cases where host memory is not needed"""
    def __init__(self, size: int):
        self._device_ptr = cuda_call(cudart.cudaMalloc(size))
        self._nbytes = size

    @property
    def device_ptr(self) -> int:
        """Device memory pointer"""
        return self._device_ptr

    @property
    def nbytes(self) -> int:
        return self._nbytes

    def free(self):
        """Explicitly free device memory"""
        if self._device_ptr is not None:
            try:
                cuda_call(cudart.cudaFree(self._device_ptr))
                self._device_ptr = None
            except Exception:
                # Log but don't raise - cleanup should be best effort
                pass

    def __str__(self):
        return f"Device:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"

    def __repr__(self):
        return self.__str__()

    def __del__(self):
        # Fallback cleanup - not guaranteed to be called
        self.free()


class PinnedHostMem:
    """Pinned host memory allocation for faster GPU transfers"""
    def __init__(self, size: int, dtype: Optional[np.dtype] = None):
        if dtype is None:
            dtype = np.dtype(np.uint8)
        else:
            dtype = np.dtype(dtype)
        nbytes = size * dtype.itemsize
        host_mem = cuda_call(cudart.cudaMallocHost(nbytes))

        self._host_ptr = host_mem
        self._host_size = size
        self._nbytes = nbytes
        self._dtype = dtype

    @property
    def array(self) -> np.ndarray:
        # Create view with proper memory ownership
        pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(self._dtype))
        host_array = np.ctypeslib.as_array(ctypes.cast(self._host_ptr, pointer_type), (self._host_size,))
        return ArrayWithOwner(host_array, self)

    @array.setter
    def array(self, data: Union[np.ndarray, bytes]):
        """Set the array data with proper bounds checking"""
        host_array = self.array  # Get the numpy array view
        if isinstance(data, np.ndarray):
            if data.size > self._host_size:
                raise ValueError(
                    f"Tried to fit an array of size {data.size} into host memory of size {self._host_size}"
                )
            np.copyto(host_array[:data.size], data.flat, casting='safe')
        else:
            assert self._dtype == np.uint8
            host_array[:self.nbytes] = np.frombuffer(data, dtype=np.uint8)


    @property
    def nbytes(self) -> int:
        return self._nbytes

    def free(self):
        """Explicitly free pinned host memory"""
        if self._host_ptr is not None:
            try:
                cuda_call(cudart.cudaFreeHost(self._host_ptr))
                self._host_ptr = None
            except Exception:
                # Log but don't raise - cleanup should be best effort
                pass

    def __str__(self):
        return f"PinnedHost:\n{self.array}\nSize:\n{self.nbytes}\n"

    def __repr__(self):
        return self.__str__()

    def __del__(self):
        # Fallback cleanup - not guaranteed to be called
        self.free()

class CudaStreamContext:
    """CUDA stream lifecycle management with context manager support"""
    def __init__(self):
        """Initialize CUDA stream"""
        self._stream = cuda_call(cudart.cudaStreamCreate())

    def __enter__(self):
        """Create CUDA stream when entering context (if not already created)"""
        if self._stream is None:
            self._stream = cuda_call(cudart.cudaStreamCreate())
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Destroy CUDA stream when exiting context"""
        if self._stream is not None:
            try:
                cuda_call(cudart.cudaStreamDestroy(self._stream))
            except Exception:
                # Silently handle cleanup failures
                pass
            self._stream = None

    @property
    def stream(self) -> cudart.cudaStream_t:
        if self._stream is None:
            raise RuntimeError("Stream not created. Use 'with' statement.")
        return self._stream

    def synchronize(self):
        """Synchronize the stream"""
        if self._stream is None:
            raise RuntimeError("Stream not created. Use 'with' statement.")
        cuda_call(cudart.cudaStreamSynchronize(self._stream))

    def free(self):
        """Explicitly free the CUDA stream"""
        if self._stream is not None:
            try:
                cuda_call(cudart.cudaStreamDestroy(self._stream))
                self._stream = None
            except Exception:
                # Log but don't raise - cleanup should be best effort
                pass

    def __del__(self):
        """Cleanup stream on destruction"""
        if hasattr(self, '_stream') and self._stream is not None:
            self.free()

    def __str__(self):
        return f"CudaStreamContext: {self._stream}"

    def __repr__(self):
        return self.__str__()


# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
# If engine uses dynamic shapes, specify a profile to find the maximum input & output size.
def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None):
    inputs = []
    outputs = []
    bindings = []
    tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
    for binding in tensor_names:
        # get_tensor_profile_shape returns (min_shape, optimal_shape, max_shape)
        # Pick out the max shape to allocate enough memory for the binding.
        shape = engine.get_tensor_shape(binding) if profile_idx is None else engine.get_tensor_profile_shape(binding, profile_idx)[-1]
        shape_valid = np.all([s >= 0 for s in shape])
        if not shape_valid and profile_idx is None:
            raise ValueError(f"Binding {binding} has dynamic shape, " +\
                "but no profile was specified.")
        size = trt.volume(shape)
        trt_type = engine.get_tensor_dtype(binding)

        # Allocate host and device buffers
        try:
            dtype = np.dtype(trt.nptype(trt_type))
            bindingMemory = HostDeviceMem(size, dtype)
        except TypeError: # no numpy support: create a byte array instead (BF16, FP8, INT4)
            size = int(size * trt_type.itemsize)
            bindingMemory = HostDeviceMem(size)

        # Append the device buffer to device bindings.
        bindings.append(int(bindingMemory.device_ptr))

        # Append to the appropriate list.
        if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
            inputs.append(bindingMemory)
        else:
            outputs.append(bindingMemory)
    return inputs, outputs, bindings


# Frees the resources allocated in allocate_buffers
def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem]):
    """
    Explicitly free CUDA memory resources.

    While __del__ methods provide automatic cleanup, they are not guaranteed to be called.
    This function provides explicit resource management for critical applications.
    """
    for inp in inputs:
        if hasattr(inp, '_device_mem') and hasattr(inp._device_mem, 'free'):
            inp._device_mem.free()
        if hasattr(inp, '_host_mem') and hasattr(inp._host_mem, 'free'):
            inp._host_mem.free()

    for out in outputs:
        if hasattr(out, '_device_mem') and hasattr(out._device_mem, 'free'):
            out._device_mem.free()
        if hasattr(out, '_host_mem') and hasattr(out._host_mem, 'free'):
            out._host_mem.free()


# Wrapper for cudaMemcpy which infers copy size and does error checking
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
    cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))

# Wrapper for cudaMemcpy which infers copy size and does error checking
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
    cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))


# Additional CUDA wrapper functions for common operations


def cuda_init():
    """Initialize CUDA driver API with error checking."""
    cuda_call(cuda.cuInit(0))


def cuda_get_device(device_id: int = 0):
    """Get CUDA device handle with error checking."""
    return cuda_call(cuda.cuDeviceGet(device_id))


# CUDA Runtime API functions (preferred over driver API when available)


def cuda_memcpy_htod(device_ptr: int, host_data: np.ndarray):
    """Copy data from host to device using CUDA runtime API with error checking.

    Note: Consider using HostDeviceMem.host setter for integrated memory management.
    """
    cuda_call(cudart.cudaMemcpy(device_ptr, host_data, host_data.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))


def _do_inference_base(inputs, outputs, stream, execute_async_func):
    # Transfer input data to the GPU.
    kind = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
    [cuda_call(cudart.cudaMemcpyAsync(inp.device_ptr, inp.host.ctypes.data, inp.nbytes, kind, stream)) for inp in inputs]
    # Run inference.
    execute_async_func()
    # Transfer predictions back from the GPU.
    kind = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
    [cuda_call(cudart.cudaMemcpyAsync(out.host.ctypes.data, out.device_ptr, out.nbytes, kind, stream)) for out in outputs]
    # Synchronize the stream
    cuda_call(cudart.cudaStreamSynchronize(stream))
    # Return only the host outputs.
    return [out.host.copy() for out in outputs]


# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, engine, bindings, inputs, outputs, stream):
    """
    Perform inference using the provided context and stream.

    Usage with context manager:
        with stream:  # Ensures proper stream lifecycle
            outputs = do_inference(context, engine, bindings, inputs, outputs, stream)
    """
    stream_handle = stream.stream
    def execute_async_func():
        context.execute_async_v3(stream_handle=stream_handle)
    # Setup context tensor address.
    num_io = engine.num_io_tensors
    for i in range(num_io):
        context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
    return _do_inference_base(inputs, outputs, stream_handle, execute_async_func)