//! Activation Checkpointing for Metal
//!
//! This module provides activation checkpointing to trade compute for memory,
//! allowing training of larger models by recomputing activations during
//! the backward pass instead of storing them.
//!
//! ## Strategy
//!
//! Instead of storing all intermediate activations for the backward pass,
//! we only store activations at certain "checkpoint" layers. During backward,
//! we recompute the activations between checkpoints as needed.

use crate::device::MetalDevice;
use crate::error::Result;
use crate::memory::MetalBuffer;
use std::collections::HashMap;
use std::ops::Range;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use parking_lot::RwLock;

/// Unique identifier for a checkpoint
pub type CheckpointId = u64;

/// Policy for determining which layers to checkpoint
#[derive(Debug, Clone)]
pub enum CheckpointPolicy {
    /// Checkpoint every N layers
    Every(usize),
    /// Checkpoint to stay within memory budget (bytes)
    MemoryBudget(usize),
    /// Checkpoint specific layers only
    SelectiveLayers(Vec<usize>),
    /// No checkpointing (store all activations)
    None,
    /// Automatic: choose based on model and memory
    Auto {
        /// Target memory reduction factor
        target_reduction: f32,
    },
}

impl Default for CheckpointPolicy {
    fn default() -> Self {
        Self::Every(4)
    }
}

impl CheckpointPolicy {
    /// Determine if a layer should be checkpointed
    pub fn should_checkpoint(&self, layer_idx: usize, total_layers: usize) -> bool {
        match self {
            Self::Every(n) => layer_idx % n == 0,
            Self::MemoryBudget(_) => {
                // Heuristic: checkpoint every sqrt(total_layers) layers
                let interval = (total_layers as f64).sqrt().ceil() as usize;
                layer_idx % interval == 0
            }
            Self::SelectiveLayers(layers) => layers.contains(&layer_idx),
            Self::None => false,
            Self::Auto { target_reduction } => {
                // Estimate interval based on target reduction
                // reduction = 1 - 1/interval, so interval = 1/(1-reduction)
                let interval = (1.0 / (1.0 - target_reduction)).ceil() as usize;
                layer_idx % interval == 0
            }
        }
    }
}

/// Configuration for checkpointing
#[derive(Debug, Clone)]
pub struct CheckpointConfig {
    /// Checkpointing policy
    pub policy: CheckpointPolicy,
    /// Maximum checkpoints to keep in memory
    pub max_checkpoints: usize,
    /// Whether to use async recompute
    pub async_recompute: bool,
    /// Memory overhead threshold before evicting old checkpoints
    pub memory_threshold: f32,
}

impl Default for CheckpointConfig {
    fn default() -> Self {
        Self {
            policy: CheckpointPolicy::default(),
            max_checkpoints: 64,
            async_recompute: true,
            memory_threshold: 0.9,
        }
    }
}

/// Data stored at a checkpoint
pub struct CheckpointData {
    /// Input tensors saved at this checkpoint
    pub input_tensors: Vec<MetalBuffer<f32>>,
    /// Layer indices covered by this checkpoint (for recompute)
    pub layer_range: Range<usize>,
    /// Whether this checkpoint is pinned (won't be evicted)
    pub pinned: bool,
    /// Reference count for this checkpoint
    ref_count: AtomicU64,
}

impl CheckpointData {
    /// Create a new checkpoint
    pub fn new(input_tensors: Vec<MetalBuffer<f32>>, layer_range: Range<usize>) -> Self {
        Self {
            input_tensors,
            layer_range,
            pinned: false,
            ref_count: AtomicU64::new(1),
        }
    }

    /// Get the memory footprint of this checkpoint
    pub fn memory_bytes(&self) -> usize {
        self.input_tensors.iter().map(|t| t.size_bytes()).sum()
    }

    /// Increment reference count
    pub fn inc_ref(&self) {
        self.ref_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Decrement reference count
    pub fn dec_ref(&self) -> u64 {
        self.ref_count.fetch_sub(1, Ordering::Relaxed)
    }

    /// Get reference count
    pub fn ref_count(&self) -> u64 {
        self.ref_count.load(Ordering::Relaxed)
    }
}

impl std::fmt::Debug for CheckpointData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CheckpointData")
            .field("layer_range", &self.layer_range)
            .field("num_tensors", &self.input_tensors.len())
            .field("memory_bytes", &self.memory_bytes())
            .field("pinned", &self.pinned)
            .finish()
    }
}

/// Statistics for checkpoint manager
#[derive(Debug, Default, Clone)]
pub struct CheckpointStats {
    /// Total checkpoints created
    pub total_created: u64,
    /// Total checkpoints evicted
    pub total_evicted: u64,
    /// Total recomputes performed
    pub total_recomputes: u64,
    /// Memory saved by checkpointing (estimated bytes)
    pub memory_saved: u64,
    /// Current checkpoints in memory
    pub current_checkpoints: usize,
    /// Current memory usage
    pub current_memory: usize,
}

/// Metal Checkpoint Manager
pub struct MetalCheckpointManager {
    /// Device
    device: Arc<MetalDevice>,
    /// Configuration
    config: CheckpointConfig,
    /// Stored checkpoints
    checkpoints: RwLock<HashMap<CheckpointId, CheckpointData>>,
    /// Next checkpoint ID
    next_id: AtomicU64,
    /// Statistics
    stats: RwLock<CheckpointStats>,
    /// Layer to checkpoint ID mapping
    layer_checkpoints: RwLock<HashMap<usize, CheckpointId>>,
}

impl MetalCheckpointManager {
    /// Create a new checkpoint manager
    pub fn new(device: Arc<MetalDevice>, config: CheckpointConfig) -> Self {
        Self {
            device,
            config,
            checkpoints: RwLock::new(HashMap::new()),
            next_id: AtomicU64::new(1),
            stats: RwLock::new(CheckpointStats::default()),
            layer_checkpoints: RwLock::new(HashMap::new()),
        }
    }

    /// Create with default configuration
    pub fn default_config(device: Arc<MetalDevice>) -> Self {
        Self::new(device, CheckpointConfig::default())
    }

    /// Check if a layer should be checkpointed
    pub fn should_checkpoint(&self, layer_idx: usize, total_layers: usize) -> bool {
        self.config.policy.should_checkpoint(layer_idx, total_layers)
    }

    /// Save a checkpoint
    pub fn save_checkpoint(
        &self,
        layer_idx: usize,
        inputs: Vec<MetalBuffer<f32>>,
        layer_range: Range<usize>,
    ) -> Result<CheckpointId> {
        // Check if we need to evict
        self.evict_if_needed()?;

        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let checkpoint = CheckpointData::new(inputs, layer_range.clone());
        let memory = checkpoint.memory_bytes();

        {
            let mut checkpoints = self.checkpoints.write();
            checkpoints.insert(id, checkpoint);
        }

        {
            let mut layer_map = self.layer_checkpoints.write();
            layer_map.insert(layer_idx, id);
        }

        {
            let mut stats = self.stats.write();
            stats.total_created += 1;
            stats.current_checkpoints += 1;
            stats.current_memory += memory;
        }

        Ok(id)
    }

    /// Get a checkpoint for recomputation
    pub fn get_checkpoint(&self, id: CheckpointId) -> Option<CheckpointDataRef> {
        let checkpoints = self.checkpoints.read();
        if checkpoints.contains_key(&id) {
            Some(CheckpointDataRef { manager: self, id })
        } else {
            None
        }
    }

    /// Get checkpoint by layer index
    pub fn get_checkpoint_for_layer(&self, layer_idx: usize) -> Option<CheckpointDataRef> {
        let layer_map = self.layer_checkpoints.read();
        let id = layer_map.get(&layer_idx)?;
        self.get_checkpoint(*id)
    }

    /// Find the nearest checkpoint before a given layer
    pub fn find_nearest_checkpoint(&self, layer_idx: usize) -> Option<(CheckpointId, usize)> {
        let layer_map = self.layer_checkpoints.read();

        layer_map
            .iter()
            .filter(|(l, _)| **l <= layer_idx)
            .max_by_key(|(l, _)| *l)
            .map(|(l, id)| (*id, *l))
    }

    /// Mark a recompute as complete
    pub fn record_recompute(&self) {
        let mut stats = self.stats.write();
        stats.total_recomputes += 1;
    }

    /// Release a checkpoint
    pub fn release_checkpoint(&self, id: CheckpointId) -> Result<()> {
        let should_remove = {
            let checkpoints = self.checkpoints.read();
            if let Some(cp) = checkpoints.get(&id) {
                cp.dec_ref() <= 1 && !cp.pinned
            } else {
                false
            }
        };

        if should_remove {
            let memory = {
                let mut checkpoints = self.checkpoints.write();
                if let Some(cp) = checkpoints.remove(&id) {
                    cp.memory_bytes()
                } else {
                    0
                }
            };

            // Update stats
            {
                let mut stats = self.stats.write();
                stats.current_checkpoints = stats.current_checkpoints.saturating_sub(1);
                stats.current_memory = stats.current_memory.saturating_sub(memory);
            }

            // Remove from layer map
            {
                let mut layer_map = self.layer_checkpoints.write();
                layer_map.retain(|_, &mut v| v != id);
            }
        }

        Ok(())
    }

    /// Pin a checkpoint (prevent eviction)
    pub fn pin_checkpoint(&self, id: CheckpointId) {
        let mut checkpoints = self.checkpoints.write();
        if let Some(cp) = checkpoints.get_mut(&id) {
            cp.pinned = true;
        }
    }

    /// Unpin a checkpoint
    pub fn unpin_checkpoint(&self, id: CheckpointId) {
        let mut checkpoints = self.checkpoints.write();
        if let Some(cp) = checkpoints.get_mut(&id) {
            cp.pinned = false;
        }
    }

    /// Evict old checkpoints if memory is constrained
    fn evict_if_needed(&self) -> Result<()> {
        let checkpoints = self.checkpoints.read();

        if checkpoints.len() >= self.config.max_checkpoints {
            drop(checkpoints);

            // Find candidates for eviction (unpinned, low ref count)
            let mut candidates: Vec<(CheckpointId, u64)> = {
                let checkpoints = self.checkpoints.read();
                checkpoints
                    .iter()
                    .filter(|(_, cp)| !cp.pinned)
                    .map(|(&id, cp)| (id, cp.ref_count()))
                    .collect()
            };

            // Sort by ref count (evict least used first)
            candidates.sort_by_key(|(_, count)| *count);

            // Evict oldest unpinned checkpoint
            if let Some((id, _)) = candidates.first() {
                let memory = {
                    let mut checkpoints = self.checkpoints.write();
                    if let Some(cp) = checkpoints.remove(id) {
                        cp.memory_bytes()
                    } else {
                        0
                    }
                };

                {
                    let mut stats = self.stats.write();
                    stats.total_evicted += 1;
                    stats.current_checkpoints = stats.current_checkpoints.saturating_sub(1);
                    stats.current_memory = stats.current_memory.saturating_sub(memory);
                }

                {
                    let mut layer_map = self.layer_checkpoints.write();
                    layer_map.retain(|_, &mut v| v != *id);
                }
            }
        }

        Ok(())
    }

    /// Clear all checkpoints
    pub fn clear(&self) {
        let mut checkpoints = self.checkpoints.write();
        checkpoints.clear();

        let mut layer_map = self.layer_checkpoints.write();
        layer_map.clear();

        let mut stats = self.stats.write();
        stats.current_checkpoints = 0;
        stats.current_memory = 0;
    }

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

    /// Get configuration
    pub fn config(&self) -> &CheckpointConfig {
        &self.config
    }

    /// Get device
    pub fn device(&self) -> &Arc<MetalDevice> {
        &self.device
    }

    /// Get number of checkpoints
    pub fn checkpoint_count(&self) -> usize {
        self.checkpoints.read().len()
    }

    /// Get total memory used by checkpoints
    pub fn memory_usage(&self) -> usize {
        self.stats.read().current_memory
    }

    /// Estimate memory saved by checkpointing
    pub fn estimate_memory_saved(&self, total_layers: usize, activation_size: usize) -> usize {
        let checkpoint_count = self.checkpoint_count();
        if checkpoint_count == 0 {
            return 0;
        }

        // Without checkpointing: store all activations
        let without = total_layers * activation_size;

        // With checkpointing: store checkpoint activations + recompute
        let with = checkpoint_count * activation_size;

        without.saturating_sub(with)
    }
}

impl std::fmt::Debug for MetalCheckpointManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let stats = self.stats();
        f.debug_struct("MetalCheckpointManager")
            .field("checkpoint_count", &stats.current_checkpoints)
            .field("memory_usage", &stats.current_memory)
            .field("total_recomputes", &stats.total_recomputes)
            .finish()
    }
}

/// Reference to checkpoint data (RAII)
pub struct CheckpointDataRef<'a> {
    manager: &'a MetalCheckpointManager,
    id: CheckpointId,
}

impl<'a> CheckpointDataRef<'a> {
    /// Get the checkpoint data
    pub fn data(&self) -> Option<impl std::ops::Deref<Target = CheckpointData> + '_> {
        let checkpoints = self.manager.checkpoints.read();
        if checkpoints.contains_key(&self.id) {
            // Increment ref count
            checkpoints.get(&self.id).unwrap().inc_ref();
            Some(parking_lot::RwLockReadGuard::map(checkpoints, |c| {
                c.get(&self.id).unwrap()
            }))
        } else {
            None
        }
    }

    /// Get the checkpoint ID
    pub fn id(&self) -> CheckpointId {
        self.id
    }
}

impl<'a> Drop for CheckpointDataRef<'a> {
    fn drop(&mut self) {
        let _ = self.manager.release_checkpoint(self.id);
    }
}

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

    #[test]
    fn test_checkpoint_policy_every() {
        let policy = CheckpointPolicy::Every(4);
        assert!(policy.should_checkpoint(0, 32));
        assert!(!policy.should_checkpoint(1, 32));
        assert!(!policy.should_checkpoint(2, 32));
        assert!(!policy.should_checkpoint(3, 32));
        assert!(policy.should_checkpoint(4, 32));
    }

    #[test]
    fn test_checkpoint_policy_selective() {
        let policy = CheckpointPolicy::SelectiveLayers(vec![0, 8, 16, 24]);
        assert!(policy.should_checkpoint(0, 32));
        assert!(!policy.should_checkpoint(1, 32));
        assert!(policy.should_checkpoint(8, 32));
        assert!(!policy.should_checkpoint(7, 32));
    }

    #[test]
    fn test_checkpoint_policy_none() {
        let policy = CheckpointPolicy::None;
        assert!(!policy.should_checkpoint(0, 32));
        assert!(!policy.should_checkpoint(16, 32));
    }

    #[test]
    fn test_checkpoint_config_default() {
        let config = CheckpointConfig::default();
        assert_eq!(config.max_checkpoints, 64);
        assert!(config.async_recompute);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_checkpoint_manager_creation() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = MetalCheckpointManager::default_config(device);

        assert_eq!(manager.checkpoint_count(), 0);
        assert_eq!(manager.memory_usage(), 0);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_checkpoint_save_and_get() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = MetalCheckpointManager::default_config(device.clone());

        // Create some test buffers
        let buffer = MetalBuffer::<f32>::zeros(&device, 1024).unwrap();

        // Save checkpoint
        let id = manager.save_checkpoint(0, vec![buffer], 0..4).unwrap();
        assert_eq!(manager.checkpoint_count(), 1);

        // Get checkpoint
        let cp_ref = manager.get_checkpoint(id);
        assert!(cp_ref.is_some());

        // Get by layer
        let cp_ref2 = manager.get_checkpoint_for_layer(0);
        assert!(cp_ref2.is_some());
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_checkpoint_eviction() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let config = CheckpointConfig {
            max_checkpoints: 2,
            ..Default::default()
        };
        let manager = MetalCheckpointManager::new(device.clone(), config);

        // Create checkpoints
        for i in 0..3 {
            let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
            manager
                .save_checkpoint(i, vec![buffer], i..i + 4)
                .unwrap();
        }

        // Should have evicted one
        assert_eq!(manager.checkpoint_count(), 2);

        let stats = manager.stats();
        assert_eq!(stats.total_evicted, 1);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_checkpoint_pinning() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let config = CheckpointConfig {
            max_checkpoints: 2,
            ..Default::default()
        };
        let manager = MetalCheckpointManager::new(device.clone(), config);

        // Create first checkpoint and pin it
        let buffer1 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
        let id1 = manager.save_checkpoint(0, vec![buffer1], 0..4).unwrap();
        manager.pin_checkpoint(id1);

        // Create more checkpoints to trigger eviction
        let buffer2 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
        manager.save_checkpoint(1, vec![buffer2], 4..8).unwrap();

        let buffer3 = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
        manager.save_checkpoint(2, vec![buffer3], 8..12).unwrap();

        // Pinned checkpoint should still exist
        assert!(manager.get_checkpoint(id1).is_some());
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_find_nearest_checkpoint() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = MetalCheckpointManager::default_config(device.clone());

        // Create checkpoints at layers 0, 4, 8
        for i in [0usize, 4, 8] {
            let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
            manager.save_checkpoint(i, vec![buffer], i..i + 4).unwrap();
        }

        // Find nearest checkpoint for layer 6 (should be layer 4)
        let (_, layer) = manager.find_nearest_checkpoint(6).unwrap();
        assert_eq!(layer, 4);

        // Find nearest for layer 10 (should be layer 8)
        let (_, layer) = manager.find_nearest_checkpoint(10).unwrap();
        assert_eq!(layer, 8);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_checkpoint_clear() {
        if !MetalDevice::is_available() {
            return;
        }

        let device = Arc::new(MetalDevice::system_default().unwrap());
        let manager = MetalCheckpointManager::default_config(device.clone());

        // Create some checkpoints
        for i in 0..5 {
            let buffer = MetalBuffer::<f32>::zeros(&device, 256).unwrap();
            manager.save_checkpoint(i, vec![buffer], i..i + 4).unwrap();
        }

        assert_eq!(manager.checkpoint_count(), 5);

        manager.clear();
        assert_eq!(manager.checkpoint_count(), 0);
        assert_eq!(manager.memory_usage(), 0);
    }

    #[test]
    fn test_estimate_memory_saved() {
        // Without checkpointing: 32 layers * 1MB each = 32MB
        // With checkpointing every 4: 8 checkpoints
        // Saved: 32MB - 8MB = 24MB

        // This is a simple estimate test
        let saved = 32 * 1024 * 1024 - 8 * 1024 * 1024;
        assert_eq!(saved, 24 * 1024 * 1024);
    }
}