//! Runtime for deserializing and managing TensorRT engines

use crate::error::{Error, Result};
use crate::logger::Logger;
use std::ffi::CStr;
use trtx_sys::*;

/// A CUDA engine containing optimized inference code
pub struct CudaEngine {
    inner: *mut TrtxCudaEngine,
}

impl CudaEngine {
    /// Get the number of I/O tensors
    pub fn get_nb_io_tensors(&self) -> Result<i32> {
        let mut count: i32 = 0;

        let result = unsafe { trtx_cuda_engine_get_nb_io_tensors(self.inner, &mut count) };

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &[]));
        }

        Ok(count)
    }

    /// Get the name of a tensor by index
    pub fn get_tensor_name(&self, index: i32) -> Result<String> {
        let mut name_ptr: *const i8 = std::ptr::null();
        let mut error_msg = [0i8; 1024];

        let result = unsafe {
            trtx_cuda_engine_get_tensor_name(
                self.inner,
                index,
                &mut name_ptr,
                error_msg.as_mut_ptr(),
                error_msg.len(),
            )
        };

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        let name = unsafe { CStr::from_ptr(name_ptr) }.to_str()?.to_string();

        Ok(name)
    }

    /// Create an execution context for inference
    pub fn create_execution_context(&self) -> Result<ExecutionContext<'_>> {
        let mut context_ptr: *mut TrtxExecutionContext = std::ptr::null_mut();
        let mut error_msg = [0i8; 1024];

        let result = unsafe {
            trtx_cuda_engine_create_execution_context(
                self.inner,
                &mut context_ptr,
                error_msg.as_mut_ptr(),
                error_msg.len(),
            )
        };

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        Ok(ExecutionContext {
            inner: context_ptr,
            _engine: std::marker::PhantomData,
        })
    }
}

impl Drop for CudaEngine {
    fn drop(&mut self) {
        if !self.inner.is_null() {
            unsafe {
                trtx_cuda_engine_destroy(self.inner);
            }
        }
    }
}

unsafe impl Send for CudaEngine {}
unsafe impl Sync for CudaEngine {}

/// Execution context for running inference
pub struct ExecutionContext<'a> {
    inner: *mut TrtxExecutionContext,
    _engine: std::marker::PhantomData<&'a CudaEngine>,
}

impl<'a> ExecutionContext<'a> {
    /// Set the address of a tensor for input or output
    ///
    /// # Safety
    ///
    /// The caller must ensure:
    /// - `data` points to valid CUDA device memory
    /// - The memory remains valid for the lifetime of inference
    /// - The memory is large enough for the tensor's size
    pub unsafe fn set_tensor_address(
        &mut self,
        name: &str,
        data: *mut std::ffi::c_void,
    ) -> Result<()> {
        let name_cstr = std::ffi::CString::new(name)?;
        let mut error_msg = [0i8; 1024];

        let result = trtx_execution_context_set_tensor_address(
            self.inner,
            name_cstr.as_ptr(),
            data,
            error_msg.as_mut_ptr(),
            error_msg.len(),
        );

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        Ok(())
    }

    /// Enqueue inference work on a CUDA stream
    ///
    /// # Safety
    ///
    /// The caller must ensure:
    /// - `cuda_stream` is a valid CUDA stream handle (or null for default stream)
    /// - All tensor addresses have been set
    /// - CUDA context is properly initialized
    pub unsafe fn enqueue_v3(&mut self, cuda_stream: *mut std::ffi::c_void) -> Result<()> {
        let mut error_msg = [0i8; 1024];

        let result = trtx_execution_context_enqueue_v3(
            self.inner,
            cuda_stream,
            error_msg.as_mut_ptr(),
            error_msg.len(),
        );

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        Ok(())
    }
}

impl Drop for ExecutionContext<'_> {
    fn drop(&mut self) {
        if !self.inner.is_null() {
            unsafe {
                trtx_execution_context_destroy(self.inner);
            }
        }
    }
}

unsafe impl Send for ExecutionContext<'_> {}

/// Runtime for deserializing engines
pub struct Runtime<'a> {
    inner: *mut TrtxRuntime,
    _logger: &'a Logger,
}

impl<'a> Runtime<'a> {
    /// Create a new runtime
    pub fn new(logger: &'a Logger) -> Result<Self> {
        let mut runtime_ptr: *mut TrtxRuntime = std::ptr::null_mut();
        let mut error_msg = [0i8; 1024];

        let result = unsafe {
            trtx_runtime_create(
                logger.as_ptr(),
                &mut runtime_ptr,
                error_msg.as_mut_ptr(),
                error_msg.len(),
            )
        };

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        Ok(Runtime {
            inner: runtime_ptr,
            _logger: logger,
        })
    }

    /// Deserialize a CUDA engine from serialized data
    pub fn deserialize_cuda_engine(&self, data: &[u8]) -> Result<CudaEngine> {
        let mut engine_ptr: *mut TrtxCudaEngine = std::ptr::null_mut();
        let mut error_msg = [0i8; 1024];

        let result = unsafe {
            trtx_runtime_deserialize_cuda_engine(
                self.inner,
                data.as_ptr() as *const std::ffi::c_void,
                data.len(),
                &mut engine_ptr,
                error_msg.as_mut_ptr(),
                error_msg.len(),
            )
        };

        if result != TRTX_SUCCESS {
            return Err(Error::from_ffi(result, &error_msg));
        }

        Ok(CudaEngine { inner: engine_ptr })
    }
}

impl Drop for Runtime<'_> {
    fn drop(&mut self) {
        if !self.inner.is_null() {
            unsafe {
                trtx_runtime_destroy(self.inner);
            }
        }
    }
}

unsafe impl Send for Runtime<'_> {}
