//! Distributed attention and KV cache management.
//!
//! This module implements distributed multi-head attention with sharded
//! KV cache for memory-efficient inference of large language models.

use distllm_shared::{DataType, KVCacheConfig};

// ============================================================================
// Distributed Attention
// ============================================================================

/// Distributed multi-head attention with tensor parallelism support.
#[derive(Debug)]
pub struct DistributedAttention {
    /// Number of query heads per shard.
    pub num_heads: usize,
    /// Number of key-value heads per shard.
    pub num_kv_heads: usize,
    /// Dimension per attention head.
    pub head_dim: usize,
    /// Tensor parallel world size.
    pub tp_world_size: usize,
    /// Whether to use flash attention.
    pub use_flash_attention: bool,
    /// RoPE theta for positional encoding.
    pub rope_theta: f64,
    /// Maximum sequence length.
    pub max_seq_len: usize,
    /// Cached cos/sin for RoPE.
    rope_cache: Option<(Vec<f32>, Vec<f32>)>,
}

impl DistributedAttention {
    /// Create a new distributed attention module.
    #[must_use]
    pub fn new(
        num_heads: usize,
        num_kv_heads: usize,
        head_dim: usize,
        tp_world_size: usize,
    ) -> Self {
        Self {
            num_heads,
            num_kv_heads,
            head_dim,
            tp_world_size,
            use_flash_attention: true,
            rope_theta: 10000.0,
            max_seq_len: 4096,
            rope_cache: None,
        }
    }

    /// Initialize RoPE cache for given sequence length.
    pub fn init_rope_cache(&mut self, seq_len: usize) {
        let half_dim = self.head_dim / 2;
        let mut cos = Vec::with_capacity(seq_len * half_dim);
        let mut sin = Vec::with_capacity(seq_len * half_dim);

        for pos in 0..seq_len {
            for i in 0..half_dim {
                let freq = 1.0 / self.rope_theta.powf(2.0 * i as f64 / self.head_dim as f64);
                let angle = pos as f64 * freq;
                cos.push(angle.cos() as f32);
                sin.push(angle.sin() as f32);
            }
        }

        self.rope_cache = Some((cos, sin));
    }

    /// Compute attention output (simulated).
    ///
    /// Returns simulated output dimensions.
    #[must_use]
    pub fn forward(&self, seq_len: usize, _layer_id: usize) -> (usize, usize, usize) {
        // Output shape: (batch, seq_len, num_heads * head_dim)
        let output_dim = self.num_heads * self.head_dim;
        (1, seq_len, output_dim)
    }

    /// Compute attention with flash attention optimization (simulated).
    #[must_use]
    pub fn flash_forward(
        &self,
        seq_len: usize,
        kv_len: usize,
        _layer_id: usize,
    ) -> FlashAttentionOutput {
        // Flash attention reduces memory from O(n^2) to O(n)
        let memory_saved_ratio = if seq_len > 512 {
            1.0 - (512.0 / seq_len as f64)
        } else {
            0.0
        };

        FlashAttentionOutput {
            output_shape: (1, seq_len, self.num_heads * self.head_dim),
            kv_len,
            memory_saved_ratio,
            num_blocks: seq_len.div_ceil(128), // Block size of 128
        }
    }

    /// Apply RoPE positional encoding (simulated).
    pub fn apply_rope(&self, _positions: &[usize]) -> bool {
        self.rope_cache.is_some()
    }

    /// Get the number of heads after tensor parallel split.
    #[must_use]
    pub fn sharded_num_heads(&self) -> usize {
        self.num_heads
    }

    /// Get the number of KV heads after tensor parallel split.
    #[must_use]
    pub fn sharded_num_kv_heads(&self) -> usize {
        self.num_kv_heads
    }

    /// Calculate GQA (Grouped Query Attention) ratio.
    #[must_use]
    pub fn gqa_ratio(&self) -> usize {
        self.num_heads / self.num_kv_heads.max(1)
    }

    /// Calculate attention memory in MB for given sequence length.
    #[must_use]
    pub fn attention_memory_mb(&self, seq_len: usize, dtype_bytes: usize) -> f64 {
        // Q, K, V, output matrices
        let qkv_size = 3 * seq_len * self.num_heads * self.head_dim * dtype_bytes;
        // Attention scores (with flash attention, this is tiled)
        let score_size = if self.use_flash_attention {
            128 * 128 * self.num_heads * dtype_bytes // Block size
        } else {
            seq_len * seq_len * self.num_heads * dtype_bytes
        };
        (qkv_size + score_size) as f64 / (1024.0 * 1024.0)
    }
}

/// Output from flash attention computation.
#[derive(Debug, Clone)]
pub struct FlashAttentionOutput {
    /// Output tensor shape.
    pub output_shape: (usize, usize, usize),
    /// KV cache length used.
    pub kv_len: usize,
    /// Memory saved compared to standard attention.
    pub memory_saved_ratio: f64,
    /// Number of blocks processed.
    pub num_blocks: usize,
}

// ============================================================================
// Flash Attention
// ============================================================================

/// Flash Attention implementation for memory-efficient attention.
#[derive(Debug)]
pub struct FlashAttention {
    /// Block size for tiling.
    pub block_size: usize,
    /// Whether to use causal masking.
    pub causal: bool,
    /// Softmax scaling factor.
    pub scale: f32,
    /// Dropout probability.
    pub dropout: f32,
}

impl FlashAttention {
    /// Create a new flash attention module.
    #[must_use]
    pub fn new(head_dim: usize) -> Self {
        Self {
            block_size: 128,
            causal: true,
            scale: 1.0 / (head_dim as f32).sqrt(),
            dropout: 0.0,
        }
    }

    /// Calculate number of blocks needed for sequence length.
    #[must_use]
    pub fn num_blocks(&self, seq_len: usize) -> usize {
        seq_len.div_ceil(self.block_size)
    }

    /// Calculate memory usage in MB.
    #[must_use]
    pub fn memory_mb(&self, seq_len: usize, num_heads: usize, dtype_bytes: usize) -> f64 {
        // Flash attention only needs O(block_size^2) per block
        let block_mem = self.block_size * self.block_size * num_heads * dtype_bytes;
        let num_blocks = self.num_blocks(seq_len);
        // We process 2 blocks at a time (Q block and KV block)
        (2 * block_mem * num_blocks) as f64 / (1024.0 * 1024.0)
    }

    /// Calculate FLOPs for attention computation.
    #[must_use]
    pub fn flops(
        &self,
        batch_size: usize,
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
    ) -> u64 {
        // QK^T: 2 * batch * heads * seq * seq * dim
        // Softmax: 5 * batch * heads * seq * seq
        // AV: 2 * batch * heads * seq * seq * dim
        let qk_flops = 2 * batch_size * num_heads * seq_len * seq_len * head_dim;
        let softmax_flops = 5 * batch_size * num_heads * seq_len * seq_len;
        let av_flops = 2 * batch_size * num_heads * seq_len * seq_len * head_dim;
        (qk_flops + softmax_flops + av_flops) as u64
    }
}

impl Default for FlashAttention {
    fn default() -> Self {
        Self::new(128)
    }
}

// ============================================================================
// KV Cache
// ============================================================================

/// Key-Value cache for efficient autoregressive generation.
#[derive(Debug)]
pub struct KVCache {
    /// Configuration.
    pub config: KVCacheConfig,
    /// Current sequence length per layer.
    pub seq_lens: Vec<usize>,
    /// Total memory allocated in MB.
    pub allocated_memory_mb: f64,
    /// Whether paged attention is enabled.
    pub paged: bool,
    /// Block tables for paged attention.
    pub block_tables: Vec<Vec<usize>>,
    /// Free block list.
    pub free_blocks: Vec<usize>,
    /// Number of allocated blocks.
    pub num_allocated_blocks: usize,
}

impl KVCache {
    /// Create a new KV cache.
    #[must_use]
    pub fn new(config: KVCacheConfig) -> Self {
        let seq_lens = vec![0; config.num_layers];
        let allocated_memory_mb = config.total_memory_gb() * 1024.0;
        let paged = config.paged_attention;

        // Initialize block management for paged attention
        let total_blocks = if config.paged_attention {
            let tokens_per_block = config.block_size;
            let total_tokens = config.max_batch_size * config.max_seq_len;
            total_tokens.div_ceil(tokens_per_block)
        } else {
            0
        };

        let free_blocks = (0..total_blocks).collect();

        Self {
            config,
            seq_lens,
            allocated_memory_mb,
            paged,
            block_tables: vec![],
            free_blocks,
            num_allocated_blocks: 0,
        }
    }

    /// Update cache for a layer with new tokens.
    pub fn update(&mut self, layer_id: usize, num_tokens: usize) {
        if layer_id < self.seq_lens.len() {
            self.seq_lens[layer_id] += num_tokens;
        }
    }

    /// Get current sequence length for a layer.
    #[must_use]
    pub fn get_seq_len(&self, layer_id: usize) -> usize {
        self.seq_lens.get(layer_id).copied().unwrap_or(0)
    }

    /// Clear cache for all layers.
    pub fn clear(&mut self) {
        for len in &mut self.seq_lens {
            *len = 0;
        }
        if self.paged {
            self.block_tables.clear();
            self.free_blocks = (0..self.free_blocks.len() + self.num_allocated_blocks).collect();
            self.num_allocated_blocks = 0;
        }
    }

    /// Allocate blocks for a new sequence.
    pub fn allocate_sequence(&mut self, seq_id: usize, initial_len: usize) -> bool {
        if !self.paged {
            return true;
        }

        let blocks_needed = initial_len.div_ceil(self.config.block_size);
        if blocks_needed > self.free_blocks.len() {
            return false;
        }

        let allocated: Vec<usize> = self.free_blocks.drain(..blocks_needed).collect();
        self.num_allocated_blocks += allocated.len();

        while self.block_tables.len() <= seq_id {
            self.block_tables.push(vec![]);
        }
        self.block_tables[seq_id] = allocated;

        true
    }

    /// Free blocks for a completed sequence.
    pub fn free_sequence(&mut self, seq_id: usize) {
        if !self.paged || seq_id >= self.block_tables.len() {
            return;
        }

        let blocks = std::mem::take(&mut self.block_tables[seq_id]);
        self.num_allocated_blocks -= blocks.len();
        self.free_blocks.extend(blocks);
    }

    /// Get memory usage ratio.
    #[must_use]
    pub fn memory_usage_ratio(&self) -> f64 {
        if self.paged {
            let total_blocks = self.free_blocks.len() + self.num_allocated_blocks;
            if total_blocks > 0 {
                self.num_allocated_blocks as f64 / total_blocks as f64
            } else {
                0.0
            }
        } else {
            let max_total: usize = self.seq_lens.len() * self.config.max_seq_len;
            let current_total: usize = self.seq_lens.iter().sum();
            if max_total > 0 {
                current_total as f64 / max_total as f64
            } else {
                0.0
            }
        }
    }

    /// Calculate per-token memory in bytes.
    #[must_use]
    pub fn per_token_memory_bytes(&self) -> usize {
        let dtype_bytes = match self.config.dtype {
            DataType::Float32 => 4,
            DataType::Float16 | DataType::BFloat16 => 2,
            DataType::Int8 | DataType::FP8 => 1,
            DataType::Int4 => 1, // Rounded up
        };
        // K and V for all layers and heads
        2 * self.config.num_layers * self.config.num_kv_heads * self.config.head_dim * dtype_bytes
    }
}

// ============================================================================
// Sharded KV Cache
// ============================================================================

/// Sharded KV cache for distributed inference.
#[derive(Debug)]
pub struct ShardedKVCache {
    /// Local cache for this shard.
    pub local_cache: KVCache,
    /// Shard ID.
    pub shard_id: usize,
    /// Total number of shards.
    pub num_shards: usize,
    /// Layer range handled by this shard.
    pub layer_range: (usize, usize),
}

impl ShardedKVCache {
    /// Create a new sharded KV cache.
    #[must_use]
    pub fn new(
        config: KVCacheConfig,
        shard_id: usize,
        num_shards: usize,
        layer_start: usize,
        layer_end: usize,
    ) -> Self {
        // Adjust config for this shard
        let shard_config = KVCacheConfig {
            num_layers: layer_end - layer_start,
            ..config
        };

        Self {
            local_cache: KVCache::new(shard_config),
            shard_id,
            num_shards,
            layer_range: (layer_start, layer_end),
        }
    }

    /// Check if this shard handles a given layer.
    #[must_use]
    pub fn handles_layer(&self, layer_id: usize) -> bool {
        layer_id >= self.layer_range.0 && layer_id < self.layer_range.1
    }

    /// Get local layer index for a global layer ID.
    #[must_use]
    pub fn local_layer_index(&self, layer_id: usize) -> Option<usize> {
        if self.handles_layer(layer_id) {
            Some(layer_id - self.layer_range.0)
        } else {
            None
        }
    }

    /// Update cache for a layer.
    pub fn update(&mut self, layer_id: usize, num_tokens: usize) {
        if let Some(local_idx) = self.local_layer_index(layer_id) {
            self.local_cache.update(local_idx, num_tokens);
        }
    }

    /// Get memory usage for this shard.
    #[must_use]
    pub fn memory_mb(&self) -> f64 {
        self.local_cache.allocated_memory_mb
    }
}

// ============================================================================
// Attention Mask
// ============================================================================

/// Attention mask types.
#[derive(Debug, Clone, Default)]
pub enum AttentionMask {
    /// No mask (full attention).
    None,
    /// Causal mask (autoregressive).
    #[default]
    Causal,
    /// Sliding window attention.
    SlidingWindow { window_size: usize },
    /// Custom mask.
    Custom { mask: Vec<Vec<bool>> },
}

impl AttentionMask {
    /// Create causal mask for given sequence length.
    #[must_use]
    pub fn causal(seq_len: usize) -> Vec<Vec<bool>> {
        (0..seq_len)
            .map(|i| (0..seq_len).map(|j| j <= i).collect())
            .collect()
    }

    /// Create sliding window mask.
    #[must_use]
    pub fn sliding_window(seq_len: usize, window_size: usize) -> Vec<Vec<bool>> {
        (0..seq_len)
            .map(|i| {
                (0..seq_len)
                    .map(|j| j <= i && i.saturating_sub(window_size - 1) <= j)
                    .collect()
            })
            .collect()
    }

    /// Get the effective attention span for a position.
    #[must_use]
    pub fn effective_span(&self, position: usize, seq_len: usize) -> usize {
        match self {
            Self::None => seq_len,
            Self::Causal => position + 1,
            Self::SlidingWindow { window_size } => (*window_size).min(position + 1),
            Self::Custom { mask } => {
                if position < mask.len() {
                    mask[position].iter().filter(|&&m| m).count()
                } else {
                    0
                }
            }
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_distributed_attention_creation() {
        let attn = DistributedAttention::new(64, 8, 128, 4);
        assert_eq!(attn.num_heads, 64);
        assert_eq!(attn.num_kv_heads, 8);
        assert_eq!(attn.gqa_ratio(), 8);
    }

    #[test]
    fn test_rope_cache_init() {
        let mut attn = DistributedAttention::new(32, 8, 128, 1);
        attn.init_rope_cache(1024);
        assert!(attn.rope_cache.is_some());
    }

    #[test]
    fn test_attention_forward() {
        let attn = DistributedAttention::new(32, 8, 128, 1);
        let (batch, seq, dim) = attn.forward(512, 0);
        assert_eq!(batch, 1);
        assert_eq!(seq, 512);
        assert_eq!(dim, 32 * 128);
    }

    #[test]
    fn test_flash_attention_output() {
        let attn = DistributedAttention::new(32, 8, 128, 1);
        let output = attn.flash_forward(2048, 2048, 0);
        assert!(output.memory_saved_ratio > 0.0);
        assert!(output.num_blocks > 0);
    }

    #[test]
    fn test_flash_attention_memory() {
        let flash = FlashAttention::new(128);
        let mem = flash.memory_mb(4096, 32, 2);
        // Should be much less than O(n^2)
        assert!(mem < 100.0);
    }

    #[test]
    fn test_flash_attention_flops() {
        let flash = FlashAttention::new(128);
        let flops = flash.flops(1, 1024, 32, 128);
        assert!(flops > 0);
    }

    #[test]
    fn test_kv_cache_creation() {
        let config = KVCacheConfig::default();
        let cache = KVCache::new(config);
        assert_eq!(cache.seq_lens.len(), 32);
    }

    #[test]
    fn test_kv_cache_update() {
        let config = KVCacheConfig::default();
        let mut cache = KVCache::new(config);
        cache.update(0, 10);
        assert_eq!(cache.get_seq_len(0), 10);
        cache.update(0, 5);
        assert_eq!(cache.get_seq_len(0), 15);
    }

    #[test]
    fn test_kv_cache_clear() {
        let config = KVCacheConfig::default();
        let mut cache = KVCache::new(config);
        cache.update(0, 100);
        cache.update(1, 50);
        cache.clear();
        assert_eq!(cache.get_seq_len(0), 0);
        assert_eq!(cache.get_seq_len(1), 0);
    }

    #[test]
    fn test_paged_attention_allocation() {
        let config = KVCacheConfig {
            paged_attention: true,
            block_size: 16,
            max_seq_len: 256,
            max_batch_size: 4,
            ..Default::default()
        };
        let mut cache = KVCache::new(config);
        assert!(cache.allocate_sequence(0, 32));
        assert_eq!(cache.num_allocated_blocks, 2);
    }

    #[test]
    fn test_paged_attention_free() {
        let config = KVCacheConfig {
            paged_attention: true,
            block_size: 16,
            max_seq_len: 256,
            max_batch_size: 4,
            ..Default::default()
        };
        let mut cache = KVCache::new(config);
        cache.allocate_sequence(0, 32);
        let initial_free = cache.free_blocks.len();
        cache.free_sequence(0);
        assert_eq!(cache.free_blocks.len(), initial_free + 2);
    }

    #[test]
    fn test_sharded_kv_cache() {
        let config = KVCacheConfig {
            num_layers: 80,
            ..Default::default()
        };
        let shard = ShardedKVCache::new(config, 0, 4, 0, 20);
        assert!(shard.handles_layer(10));
        assert!(!shard.handles_layer(30));
    }

    #[test]
    fn test_sharded_cache_local_index() {
        let config = KVCacheConfig::default();
        let shard = ShardedKVCache::new(config, 1, 4, 20, 40);
        assert_eq!(shard.local_layer_index(25), Some(5));
        assert_eq!(shard.local_layer_index(10), None);
    }

    #[test]
    fn test_causal_mask() {
        let mask = AttentionMask::causal(4);
        assert!(mask[0][0]); // Can attend to position 0
        assert!(!mask[0][1]); // Cannot attend to future
        assert!(mask[3][0]); // Can attend to all previous
        assert!(mask[3][3]); // Can attend to self
    }

    #[test]
    fn test_sliding_window_mask() {
        let mask = AttentionMask::sliding_window(10, 3);
        assert!(mask[5][5]); // Self
        assert!(mask[5][4]); // Previous
        assert!(mask[5][3]); // 2 positions back
        assert!(!mask[5][2]); // 3 positions back (outside window)
    }

    #[test]
    fn test_effective_span() {
        let causal = AttentionMask::Causal;
        assert_eq!(causal.effective_span(5, 10), 6);

        let window = AttentionMask::SlidingWindow { window_size: 3 };
        assert_eq!(window.effective_span(5, 10), 3);
        assert_eq!(window.effective_span(1, 10), 2);
    }
}