//! SRAM management for Flash Attention tiling

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

/// SRAM allocation block
#[derive(Debug, Clone)]
pub struct SRAMBlock {
    /// Block ID for tracking
    pub id: u64,
    /// Starting offset in shared memory
    pub offset: usize,
    /// Size in bytes
    pub size: usize,
    /// Whether the block is currently in use
    pub in_use: bool,
    /// Block type (Q, K, V, or scratch)
    pub block_type: SRAMBlockType,
}

/// Type of SRAM block
#[derive(Debug, Clone, PartialEq)]
pub enum SRAMBlockType {
    /// Query block
    Query,
    /// Key block
    Key,
    /// Value block
    Value,
    /// Scratch space for intermediate computations
    Scratch,
    /// Attention scores
    Scores,
}

/// SRAM usage statistics
#[derive(Debug, Clone)]
pub struct SRAMStats {
    /// Total SRAM capacity in bytes
    pub total_capacity: usize,
    /// Currently allocated SRAM in bytes
    pub allocated: usize,
    /// Peak SRAM usage in bytes
    pub peak_usage: usize,
    /// Number of allocations
    pub allocation_count: u64,
    /// SRAM utilization efficiency (0.0 - 1.0)
    pub efficiency: f32,
    /// Number of allocation failures
    pub allocation_failures: u64,
}

/// SRAM memory manager for Flash Attention tiling
pub struct SRAMManager {
    /// Configuration
    config: FlashAttentionConfig,
    /// CUDA backend
    cuda_backend: Arc<CudaBackend>,
    /// Allocated blocks
    blocks: Arc<RwLock<Vec<SRAMBlock>>>,
    /// Statistics
    stats: Arc<RwLock<SRAMStats>>,
    /// Next block ID
    next_block_id: Arc<RwLock<u64>>,
    /// Total SRAM capacity per block
    total_capacity: usize,
}

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

        // Get device shared memory info
        let device_info = cuda_backend.get_device_info()
        .map_err(|e| FlashError::memory(format!("Failed to get device info: {e}")))?;

        let total_capacity = device_info.shared_memory_per_block as usize;

        // Validate configuration fits in SRAM
        let required_sram = Self::calculate_required_sram(config)?;
        if required_sram > total_capacity {
            return Err(FlashError::config(format!(
            "Configuration requires {required_sram} bytes of SRAM but only {total_capacity} bytes available"
            )));
        }

        let stats = SRAMStats {
            total_capacity,
            allocated: 0,
            peak_usage: 0,
            allocation_count: 0,
            efficiency: 0.0,
            allocation_failures: 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)),
            total_capacity,
        })
    }

    /// Calculate required SRAM for the configuration
    pub fn calculate_required_sram(config: &FlashAttentionConfig) -> FlashResult<usize> {
        let element_size = match config.precision {
            crate::config::PrecisionMode::FP32 => 4,
            crate::config::PrecisionMode::FP16 | crate::config::PrecisionMode::BF16 => 2,
            crate::config::PrecisionMode::FP8E4M3 { .. } | crate::config::PrecisionMode::FP8E5M2 { .. } => 1,
            crate::config::PrecisionMode::Mixed { storage_precision, .. } => {
                match storage_precision {
                    crate::config::Precision::FP32 => 4,
                    crate::config::Precision::FP16 | crate::config::Precision::BF16 => 2,
                    crate::config::Precision::INT8
                    | crate::config::Precision::FP8E4M3
                    | crate::config::Precision::FP8E5M2 => 1,
                }
            }
        };

        // Memory for Q block
        let q_memory = config.block_size_q * config.head_dim * element_size;

        // Memory for K block
        let k_memory = config.block_size_kv * config.head_dim * element_size;

        // Memory for V block
        let v_memory = config.block_size_kv * config.head_dim * element_size;

        // Memory for attention scores (FP32 for numerical stability)
        let scores_memory = config.block_size_q * config.block_size_kv * 4;

        // Scratch memory for intermediate computations
        let scratch_memory = config.block_size_q * config.head_dim * 4; // FP32 accumulation

        // Add alignment padding (16-byte alignment)
        let alignment = 16;
        let total = q_memory + k_memory + v_memory + scores_memory + scratch_memory;
        let aligned_total = total.div_ceil(alignment) * alignment;

        debug!("SRAM requirements: Q={}, K={}, V={}, scores={}, scratch={}, total={}",
        q_memory, k_memory, v_memory, scores_memory, scratch_memory, aligned_total);

        Ok(aligned_total)
    }

    /// Allocate SRAM block
    pub fn allocate_block(
    &self,
    size: usize,
    block_type: SRAMBlockType,
    alignment: usize,
    ) -> FlashResult<SRAMBlock> {
        let mut blocks = self.blocks.write();
        let mut stats = self.stats.write();

        // Find suitable offset with alignment
        let offset = self.find_suitable_offset(&blocks, size, alignment)?;

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

        let block = SRAMBlock {
            id: block_id,
            offset,
            size,
            in_use: true,
            block_type,
        };

        // Update statistics
        stats.allocated += size;
        stats.peak_usage = stats.peak_usage.max(stats.allocated);
        stats.allocation_count += 1;
        stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);

        blocks.push(block.clone());

        debug!("Allocated SRAM block: id={}, offset={}, size={}, type={:?}",
        block.id, block.offset, block.size, block.block_type);

        Ok(block)
    }

    /// Deallocate SRAM block
    pub fn deallocate_block(&self, block_id: u64) -> FlashResult<()> {
        let mut blocks = self.blocks.write();
        let mut stats = self.stats.write();

        if let Some(pos) = blocks.iter().position(|b| b.id == block_id) {
            let block = blocks.remove(pos);
            stats.allocated -= block.size;
            stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);

            debug!("Deallocated SRAM block: id={}, size={}", block_id, block.size);
            Ok(())
            } else {
                Err(FlashError::memory(format!("SRAM block {block_id} not found")))
            }
        }

        /// Get standard block layout for Flash Attention
        pub fn get_standard_layout(&self) -> FlashResult<SRAMLayout> {
            let element_size = match self.config.precision {
                crate::config::PrecisionMode::FP32 => 4,
                crate::config::PrecisionMode::FP16 | crate::config::PrecisionMode::BF16 => 2,
                crate::config::PrecisionMode::FP8E4M3 { .. } | crate::config::PrecisionMode::FP8E5M2 { .. } => 1,
                crate::config::PrecisionMode::Mixed { storage_precision, .. } => {
                    match storage_precision {
                        crate::config::Precision::FP32 => 4,
                        crate::config::Precision::FP16 | crate::config::Precision::BF16 => 2,
                        crate::config::Precision::INT8
                        | crate::config::Precision::FP8E4M3
                        | crate::config::Precision::FP8E5M2 => 1,
                    }
                }
            };

            let alignment = 16; // 16-byte alignment for vectorization
            let mut offset = 0;

            // Q block
            let q_size = self.config.block_size_q * self.config.head_dim * element_size;
            let q_offset = offset;
            offset = Self::align_offset(offset + q_size, alignment);

            // K block
            let k_size = self.config.block_size_kv * self.config.head_dim * element_size;
            let k_offset = offset;
            offset = Self::align_offset(offset + k_size, alignment);

            // V block
            let v_size = self.config.block_size_kv * self.config.head_dim * element_size;
            let v_offset = offset;
            offset = Self::align_offset(offset + v_size, alignment);

            // Scores block (FP32)
            let scores_size = self.config.block_size_q * self.config.block_size_kv * 4;
            let scores_offset = offset;
            offset = Self::align_offset(offset + scores_size, alignment);

            // Scratch block (FP32)
            let scratch_size = self.config.block_size_q * self.config.head_dim * 4;
            let scratch_offset = offset;
            offset = Self::align_offset(offset + scratch_size, alignment);

            if offset > self.total_capacity {
                return Err(FlashError::memory(format!(
                "SRAM layout requires {} bytes but only {} available",
                offset, self.total_capacity
                )));
            }

            Ok(SRAMLayout {
                q_block: SRAMBlockInfo { offset: q_offset, size: q_size },
                k_block: SRAMBlockInfo { offset: k_offset, size: k_size },
                v_block: SRAMBlockInfo { offset: v_offset, size: v_size },
                scores_block: SRAMBlockInfo { offset: scores_offset, size: scores_size },
                scratch_block: SRAMBlockInfo { offset: scratch_offset, size: scratch_size },
                total_size: offset,
            })
        }

        /// Estimate SRAM requirements for a sequence length
        pub fn estimate_requirements(&self, seq_len: usize) -> FlashResult<usize> {
            // Calculate how many blocks we need
            let _q_blocks = seq_len.div_ceil(self.config.block_size_q);
            let _kv_blocks = seq_len.div_ceil(self.config.block_size_kv);

            // SRAM is reused across blocks, so we only need space for one set of blocks
            let standard_layout = self.get_standard_layout()?;

            // Add some overhead for runtime management
            let overhead = standard_layout.total_size / 10; // 10% overhead

            Ok(standard_layout.total_size + overhead)
        }

        /// Get current SRAM statistics
        pub fn stats(&self) -> SRAMStats {
            self.stats.read().clone()
        }

        /// Get SRAM efficiency (0.0 - 1.0)
        pub fn get_efficiency(&self) -> f32 {
            self.stats.read().efficiency
        }

        /// Optimize SRAM allocation
        pub fn optimize_allocation(&self) -> FlashResult<()> {
            debug!("Optimizing SRAM allocation");

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

            // Sort blocks by offset to identify gaps
            blocks.sort_by_key(|b| b.offset);

            // Defragment by moving blocks to eliminate gaps
            let mut new_offset = 0;
            for block in blocks.iter_mut() {
                if !block.in_use {
                    continue;
                }

                if block.offset != new_offset {
                    debug!("Moving SRAM block {} from offset {} to {}",
                    block.id, block.offset, new_offset);
                    block.offset = new_offset;
                }

                new_offset = Self::align_offset(new_offset + block.size, 16);
            }

            // Remove unused blocks
            blocks.retain(|b| b.in_use);

            // Update efficiency
            let mut stats = self.stats.write();
            stats.efficiency = (stats.allocated as f32) / (self.total_capacity as f32);

            info!("SRAM optimization complete, efficiency: {:.2}", stats.efficiency);
            Ok(())
        }

        /// Calculate shared memory per block for kernel launch
        pub fn calculate_shared_memory_per_block(&self) -> FlashResult<usize> {
            let layout = self.get_standard_layout()?;
            Ok(layout.total_size)
        }

        /// Find suitable offset for allocation
        fn find_suitable_offset(
        &self,
        blocks: &[SRAMBlock],
        size: usize,
        alignment: usize,
        ) -> FlashResult<usize> {
            // Sort blocks by offset
            let mut sorted_blocks: Vec<_> = blocks.iter().filter(|b| b.in_use).collect();
            sorted_blocks.sort_by_key(|b| b.offset);

            let mut offset = 0;

            // Check gaps between blocks
            for block in &sorted_blocks {
                let aligned_offset = Self::align_offset(offset, alignment);
                if aligned_offset + size <= block.offset {
                    return Ok(aligned_offset);
                }
                offset = block.offset + block.size;
            }

            // Check if we can fit at the end
            let aligned_offset = Self::align_offset(offset, alignment);
            if aligned_offset + size <= self.total_capacity {
                Ok(aligned_offset)
                } else {
                    Err(FlashError::memory(format!(
                    "Cannot allocate {} bytes of SRAM (available: {})",
                    size, self.total_capacity - offset
                    )))
                }
            }

            /// Align offset to the specified alignment
            fn align_offset(offset: usize, alignment: usize) -> usize {
                offset.div_ceil(alignment) * alignment
            }
        }

        /// SRAM layout information
        #[derive(Debug, Clone)]
        pub struct SRAMLayout {
            /// Query block layout
            pub q_block: SRAMBlockInfo,
            /// Key block layout
            pub k_block: SRAMBlockInfo,
            /// Value block layout
            pub v_block: SRAMBlockInfo,
            /// Attention scores block layout
            pub scores_block: SRAMBlockInfo,
            /// Scratch space block layout
            pub scratch_block: SRAMBlockInfo,
            /// Total SRAM usage
            pub total_size: usize,
        }

        /// SRAM block information
        #[derive(Debug, Clone)]
        pub struct SRAMBlockInfo {
            /// Offset within shared memory
            pub offset: usize,
            /// Size in bytes
            pub size: usize,
        }

        impl SRAMLayout {
            /// Get the CUDA kernel parameters for this layout
            pub fn to_kernel_params(&self) -> SRAMKernelParams {
                SRAMKernelParams {
                    q_offset: self.q_block.offset as u32,
                    k_offset: self.k_block.offset as u32,
                    v_offset: self.v_block.offset as u32,
                    scores_offset: self.scores_block.offset as u32,
                    scratch_offset: self.scratch_block.offset as u32,
                    total_sram: self.total_size as u32,
                }
            }
        }

        /// SRAM kernel parameters for CUDA kernels
        #[derive(Debug, Clone)]
        pub struct SRAMKernelParams {
            pub q_offset: u32,
            pub k_offset: u32,
            pub v_offset: u32,
            pub scores_offset: u32,
            pub scratch_offset: u32,
            pub total_sram: u32,
        }

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

            #[test]
            fn test_sram_requirements_calculation() {
                let config = FlashAttentionConfig::new(32, 128);
                let required = SRAMManager::calculate_required_sram(&config).unwrap();

                // Should be reasonable amount (less than 64KB for most configurations)
                assert!(required > 0);
                assert!(required < 65_536);
            }

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

                if let Ok(cuda_backend) = CudaBackend::new(rtx_runtime::DeviceId(0)) {
                    let sram_manager = SRAMManager::new(&config, &Arc::new(cuda_backend)).unwrap();
                    let layout = sram_manager.get_standard_layout().unwrap();

                    assert!(layout.total_size > 0);
                    assert!(layout.q_block.size > 0);
                    assert!(layout.k_block.size > 0);
                    assert!(layout.v_block.size > 0);
                    assert!(layout.scores_block.size > 0);
                    assert!(layout.scratch_block.size > 0);
                }
            }

            #[test]
            fn test_alignment() {
                assert_eq!(SRAMManager::align_offset(10, 16), 16);
                assert_eq!(SRAMManager::align_offset(16, 16), 16);
                assert_eq!(SRAMManager::align_offset(17, 16), 32);
            }
        }