vecboost 0.2.0

// Copyright (c) 2025-2026 Kirky.X
//
// Licensed under the MIT License
// See LICENSE file in the project root for full license information.

use log::{debug, info, warn};
use std::sync::atomic::{AtomicU64, Ordering};

use super::config::CudaPoolConfig;

/// CUDA 内存池
///
/// 使用 CUDA API 管理底层 GPU 内存分配
#[cfg(feature = "cuda")]
pub struct CudaMemoryPool {
    /// 设备 ID
    device_id: i32,
    /// 已分配内存（字节）
    allocated_memory: AtomicU64,
    /// 最大内存（字节）
    max_memory: u64,
    /// 配置
    config: CudaPoolConfig,
}

#[cfg(feature = "cuda")]
impl CudaMemoryPool {
    /// 创建新的 CUDA 内存池
    pub fn new(device_id: i32, config: CudaPoolConfig) -> Result<Self, String> {
        let max_memory = (config.max_memory_mb * 1024 * 1024) as u64;

        info!(
            "Creating CudaMemoryPool for device {} with max_memory={}MB",
            device_id, config.max_memory_mb
        );

        // TODO: 使用 cudarc 或 CUDA API 初始化内存池
        // 这里简化实现，实际应该使用 CUDA 的内存池 API

        Ok(Self {
            device_id,
            allocated_memory: AtomicU64::new(0),
            max_memory,
            config,
        })
    }

    /// 分配内存
    pub fn allocate(&mut self, size: usize) -> Result<CudaMemoryPtr, String> {
        let size_u64 = size as u64;

        // 检查是否有足够内存
        let current_allocated = self.allocated_memory.load(Ordering::Relaxed);
        let available = self.max_memory.saturating_sub(current_allocated);

        if size_u64 > available {
            return Err(format!(
                "Insufficient CUDA memory: need {}MB, available {}MB",
                size_u64 / 1024 / 1024,
                available / 1024 / 1024
            ));
        }

        // 分配内存
        self.allocated_memory.fetch_add(size_u64, Ordering::Relaxed);

        debug!(
            "Allocated {}MB CUDA memory on device {}, total allocated: {}MB",
            size_u64 / 1024 / 1024,
            self.device_id,
            self.allocated_memory.load(Ordering::Relaxed) / 1024 / 1024
        );

        // 注意：这是占位符实现
        // 实际的 CUDA 内存分配需要 CudaDevice 句柄
        warn!(
            "CUDA memory pool is using placeholder implementation. Actual memory allocation requires CudaDevice handle."
        );

        Ok(CudaMemoryPtr {
            device_id: self.device_id,
            size,
            ptr: 0, // 占位符 - 实际应该是真实的 CUDA 指针
        })
    }

    /// 释放内存
    pub fn deallocate(&mut self, ptr: CudaMemoryPtr) {
        if ptr.size == 0 {
            warn!("Attempted to deallocate CUDA memory with size 0");
            return;
        }

        self.allocated_memory
            .fetch_sub(ptr.size as u64, Ordering::Relaxed);

        debug!(
            "Deallocated {}MB CUDA memory on device {}, total allocated: {}MB",
            ptr.size / 1024 / 1024,
            self.device_id,
            self.allocated_memory.load(Ordering::Relaxed) / 1024 / 1024
        );

        // 注意：这是占位符实现
        // 实际的 CUDA 内存释放需要 CudaDevice 句柄
        warn!(
            "CUDA memory pool is using placeholder implementation. Actual memory deallocation requires CudaDevice handle."
        );
    }

    /// 获取内存使用情况
    pub fn get_memory_usage(&self) -> (u64, u64) {
        let used = self.allocated_memory.load(Ordering::Relaxed);
        (used, self.max_memory)
    }

    /// 获取内存使用率
    pub fn get_memory_usage_percent(&self) -> f64 {
        let used = self.allocated_memory.load(Ordering::Relaxed) as f64;
        let total = self.max_memory as f64;
        (used / total) * 100.0
    }

    /// 清空池
    pub fn clear(&mut self) {
        info!("Clearing CUDA memory pool on device {}...", self.device_id);
        self.allocated_memory.store(0, Ordering::Relaxed);
        info!("CUDA memory pool cleared");
    }
}

/// CUDA 内存指针
#[cfg(feature = "cuda")]
#[derive(Debug, Clone)]
pub struct CudaMemoryPtr {
    /// 设备 ID
    pub device_id: i32,
    /// 大小（字节）
    pub size: usize,
    /// 指针（简化实现）
    pub ptr: usize,
}

#[cfg(feature = "cuda")]
impl Drop for CudaMemoryPtr {
    fn drop(&mut self) {
        debug!("Dropping CUDA memory ptr on device {}", self.device_id);
        // TODO: 使用 CUDA API 释放内存
    }
}

#[cfg(not(feature = "cuda"))]
/// CUDA 内存池（非 CUDA 版本的占位符）
pub struct CudaMemoryPool {
    _device_id: i32,
    _config: CudaPoolConfig,
}

#[cfg(not(feature = "cuda"))]
impl CudaMemoryPool {
    /// 创建新的 CUDA 内存池（非 CUDA 版本）
    pub fn new(_device_id: i32, _config: CudaPoolConfig) -> Result<Self, String> {
        warn!("CUDA feature not enabled, CudaMemoryPool will be no-op");
        Ok(Self {
            _device_id: 0,
            _config,
        })
    }

    /// 分配内存（非 CUDA 版本，返回错误）
    pub fn allocate(&mut self, _size: usize) -> Result<CudaMemoryPtr, String> {
        Err("CUDA feature not enabled".to_string())
    }

    /// 释放内存（非 CUDA 版本）
    pub fn deallocate(&mut self, _ptr: CudaMemoryPtr) {
        // No-op
    }

    /// 获取内存使用情况（非 CUDA 版本）
    pub fn get_memory_usage(&self) -> (u64, u64) {
        (0, 0)
    }

    /// 获取内存使用率（非 CUDA 版本）
    pub fn get_memory_usage_percent(&self) -> f64 {
        0.0
    }

    /// 清空池（非 CUDA 版本）
    pub fn clear(&mut self) {
        // No-op
    }
}

#[cfg(not(feature = "cuda"))]
/// CUDA 内存指针（非 CUDA 版本的占位符）
#[derive(Debug, Clone)]
pub struct CudaMemoryPtr {
    _size: usize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[cfg(feature = "cuda")]
    fn test_cuda_pool_creation() {
        let config = CudaPoolConfig::default();
        let pool = CudaMemoryPool::new(0, config);

        assert!(pool.is_ok());

        let pool = pool.unwrap();
        let (used, total) = pool.get_memory_usage();
        assert_eq!(used, 0);
        assert!(total > 0);
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_cuda_allocate_deallocate() {
        let config = CudaPoolConfig {
            max_memory_mb: 1024,
            ..Default::default()
        };

        let mut pool = CudaMemoryPool::new(0, config).unwrap();

        // 分配 512MB
        let ptr = pool.allocate(512 * 1024 * 1024);
        assert!(ptr.is_ok());

        let (used, _) = pool.get_memory_usage();
        assert_eq!(used, 512 * 1024 * 1024);

        // 释放内存
        let ptr = ptr.unwrap();
        pool.deallocate(ptr);

        let (used, _) = pool.get_memory_usage();
        assert_eq!(used, 0);
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_cuda_insufficient_memory() {
        let config = CudaPoolConfig {
            max_memory_mb: 512,
            ..Default::default()
        };

        let mut pool = CudaMemoryPool::new(0, config).unwrap();

        // 分配 512MB
        let _ptr1 = pool.allocate(512 * 1024 * 1024).unwrap();

        // 尝试再分配 512MB，应该失败
        let ptr2 = pool.allocate(512 * 1024 * 1024);
        assert!(ptr2.is_err());
    }

    #[test]
    #[cfg(not(feature = "cuda"))]
    fn test_cuda_pool_no_cuda() {
        let config = CudaPoolConfig::default();
        let pool = CudaMemoryPool::new(0, config);

        assert!(pool.is_ok());

        let mut pool = pool.unwrap();

        // 尝试分配应该失败
        let ptr = pool.allocate(1024);
        assert!(ptr.is_err());
    }

    #[test]
    #[cfg(not(feature = "cuda"))]
    fn test_cuda_pool_no_cuda_memory_usage() {
        let config = CudaPoolConfig {
            max_memory_mb: 2048,
            ..Default::default()
        };
        let pool = CudaMemoryPool::new(0, config).unwrap();

        let (used, total) = pool.get_memory_usage();
        assert_eq!(used, 0);
        assert_eq!(total, 0);

        let percent = pool.get_memory_usage_percent();
        assert_eq!(percent, 0.0);
    }

    #[test]
    #[cfg(not(feature = "cuda"))]
    fn test_cuda_pool_no_cuda_clear_noop() {
        let config = CudaPoolConfig::default();
        let mut pool = CudaMemoryPool::new(1, config).unwrap();

        pool.clear();
        pool.clear();

        let (used, total) = pool.get_memory_usage();
        assert_eq!(used, 0);
        assert_eq!(total, 0);
    }

    #[test]
    #[cfg(not(feature = "cuda"))]
    fn test_cuda_pool_no_cuda_new_with_different_device_ids() {
        for device_id in [0, 1, -1, 42] {
            let config = CudaPoolConfig::default();
            let pool = CudaMemoryPool::new(device_id, config);
            assert!(
                pool.is_ok(),
                "Failed to create pool for device {}",
                device_id
            );
        }
    }
}