//! Block memory manager for Flash Attention

use crate::{
    config::FlashAttentionConfig,
    error::{FlashError, FlashResult},
};
use rtx_runtime::CudaBackend;
use std::sync::Arc;
use parking_lot::RwLock;
use tracing::{debug, info};

/// Memory block for GPU allocation
#[derive(Debug, Clone)]
pub struct MemoryBlock {
    /// Block ID
    pub id: u64,
    /// GPU memory pointer
    pub ptr: usize,
    /// Size in bytes
    pub size: usize,
    /// Whether the block is currently allocated
    pub allocated: bool,
}

/// Block manager statistics
#[derive(Debug, Clone)]
pub struct BlockStats {
    /// Total number of blocks
    pub total_blocks: usize,
    /// Number of allocated blocks
    pub allocated_blocks: usize,
    /// Total managed memory in bytes
    pub total_memory: usize,
    /// Currently allocated memory in bytes
    pub allocated_memory: usize,
    /// Fragmentation ratio (0.0 - 1.0)
    pub fragmentation_ratio: f32,
}

/// Block memory manager
pub struct BlockManager {
    /// Configuration
    config: FlashAttentionConfig,
    /// CUDA backend
    cuda_backend: Arc<CudaBackend>,
    /// Managed blocks
    blocks: Arc<RwLock<Vec<MemoryBlock>>>,
    /// Statistics
    stats: Arc<RwLock<BlockStats>>,
    /// Next block ID
    next_block_id: Arc<RwLock<u64>>,
}

impl BlockManager {
    /// Create a new block manager
    pub fn new(config: &FlashAttentionConfig, cuda_backend: &Arc<CudaBackend>) -> FlashResult<Self> {
        info!("Initializing block manager");

        let stats = BlockStats {
            total_blocks: 0,
            allocated_blocks: 0,
            total_memory: 0,
            allocated_memory: 0,
            fragmentation_ratio: 0.0,
        };

        Ok(Self {
            config: config.clone(),
            cuda_backend: cuda_backend.clone(),
            blocks: Arc::new(RwLock::new(Vec::new())),
            stats: Arc::new(RwLock::new(stats)),
            next_block_id: Arc::new(RwLock::new(0)),
        })
    }

    /// Allocate a memory block
    pub fn allocate_block(&self, size: usize) -> FlashResult<MemoryBlock> {
        debug!("Allocating memory block of size {}", size);

        // Try to find an existing free block
        if let Some(block) = self.find_free_block(size)? {
            self.mark_block_allocated(block.id)?;
            return Ok(block);
        }

        // Allocate new block
        let cuda_slice: cudarc::driver::CudaSlice<u8> = self.cuda_backend.allocate_memory(size)
        .map_err(|e| FlashError::memory(format!("Failed to allocate GPU memory: {e}")))?;
        let ptr = &cuda_slice as *const _ as usize;

        let block_id = {
            let mut next_id = self.next_block_id.write();
            let id = *next_id;
            *next_id += 1;
            id
        };

        let block = MemoryBlock {
            id: block_id,
            ptr,
            size,
            allocated: true,
        };

        // Add to managed blocks
        {
            let mut blocks = self.blocks.write();
            blocks.push(block.clone());
        }

        // Update statistics
        {
            let mut stats = self.stats.write();
            stats.total_blocks += 1;
            stats.allocated_blocks += 1;
            stats.total_memory += size;
            stats.allocated_memory += size;
        }

        debug!("Allocated memory block: id={}, ptr=0x{:x}, size={}",
        block.id, block.ptr, block.size);

        Ok(block)
    }

    /// Deallocate a memory block
    pub fn deallocate_block(&self, block_id: u64) -> FlashResult<()> {
        debug!("Deallocating memory block {}", block_id);

        let mut blocks = self.blocks.write();
        if let Some(block) = blocks.iter_mut().find(|b| b.id == block_id) {
            if !block.allocated {
                return Err(FlashError::memory(format!("Block {block_id} is already deallocated")));
            }

            block.allocated = false;

            // Update statistics
            let mut stats = self.stats.write();
            stats.allocated_blocks -= 1;
            stats.allocated_memory -= block.size;
            stats.fragmentation_ratio = self.calculate_fragmentation(&blocks);

            debug!("Deallocated memory block {}", block_id);
            Ok(())
            } else {
                Err(FlashError::memory(format!("Block {block_id} not found")))
            }
        }

        /// Get block manager statistics
        pub fn stats(&self) -> BlockStats {
            self.stats.read().clone()
        }

        /// Defragment memory by consolidating free blocks
        pub fn defragment(&self) -> FlashResult<()> {
            info!("Starting memory defragmentation");

            let mut blocks = self.blocks.write();

            // Remove deallocated blocks
            let mut removed_memory = 0;
            blocks.retain(|block| {
                if !block.allocated {
                    removed_memory += block.size;
                    // Note: In a real implementation, we would call cuda_backend.deallocate_memory(block.ptr)
                    false
                    } else {
                        true
                    }
                });

                // Update statistics
                {
                    let mut stats = self.stats.write();
                    stats.total_blocks = blocks.len();
                    stats.total_memory -= removed_memory;
                    stats.fragmentation_ratio = self.calculate_fragmentation(&blocks);
                }

                info!("Defragmentation complete: removed {} bytes", removed_memory);
                Ok(())
            }

            /// Find a free block of at least the requested size
            fn find_free_block(&self, size: usize) -> FlashResult<Option<MemoryBlock>> {
                let blocks = self.blocks.read();
                for block in blocks.iter() {
                    if !block.allocated && block.size >= size {
                        return Ok(Some(block.clone()));
                    }
                }
                Ok(None)
            }

            /// Mark a block as allocated
            fn mark_block_allocated(&self, block_id: u64) -> FlashResult<()> {
                let mut blocks = self.blocks.write();
                if let Some(block) = blocks.iter_mut().find(|b| b.id == block_id) {
                    block.allocated = true;

                    // Update statistics
                    let mut stats = self.stats.write();
                    stats.allocated_blocks += 1;
                    stats.allocated_memory += block.size;

                    Ok(())
                    } else {
                        Err(FlashError::memory(format!("Block {block_id} not found")))
                    }
                }

                /// Calculate memory fragmentation ratio
                fn calculate_fragmentation(&self, blocks: &[MemoryBlock]) -> f32 {
                    if blocks.is_empty() {
                        return 0.0;
                    }

                    let free_blocks = blocks.iter().filter(|b| !b.allocated).count();
                    let total_blocks = blocks.len();

                    if total_blocks == 0 {
                        0.0
                        } else {
                            (free_blocks as f32) / (total_blocks as f32)
                        }
                    }
                }

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

                    #[test]
                    fn test_block_manager_creation() {
                        let config = FlashAttentionConfig::new(8, 64);

                        if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
                            let block_manager = BlockManager::new(&config, &Arc::new(cuda_backend)).unwrap();
                            let stats = block_manager.stats();

                            assert_eq!(stats.total_blocks, 0);
                            assert_eq!(stats.allocated_blocks, 0);
                            assert_eq!(stats.total_memory, 0);
                        }
                    }

                    #[test]
                    fn test_fragmentation_calculation() {
                        let config = FlashAttentionConfig::new(8, 64);

                        if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
                            let block_manager = BlockManager::new(&config, &Arc::new(cuda_backend)).unwrap();

                            // Test with empty blocks
                            let empty_blocks = vec![];
                            let fragmentation = block_manager.calculate_fragmentation(&empty_blocks);
                            assert_eq!(fragmentation, 0.0);

                            // Test with mixed blocks
                            let mixed_blocks = vec![
                            MemoryBlock { id: 1, ptr: 0x1000, size: 1024, allocated: true },
                            MemoryBlock { id: 2, ptr: 0x2000, size: 1024, allocated: false },
                            MemoryBlock { id: 3, ptr: 0x3000, size: 1024, allocated: true },
                            MemoryBlock { id: 4, ptr: 0x4000, size: 1024, allocated: false },
                            ];
                            let fragmentation = block_manager.calculate_fragmentation(&mixed_blocks);
                            assert_eq!(fragmentation, 0.5); // 2 free out of 4 total
                        }
                    }
                }