lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// CXL Memory Tiering
// Multi-tier memory hierarchy with automatic promotion/demotion.

use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use lazy_static::lazy_static;
use spin::Mutex;

/// CXL memory tier levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum CxlTier {
    /// Local DRAM (fastest, most expensive)
    LocalDram = 0,
    /// CXL-attached memory on same socket (near)
    CxlNear = 1,
    /// CXL-attached memory on remote socket (far)
    CxlFar = 2,
    /// Traditional storage (slowest, cheapest)
    Storage = 3,
}

impl CxlTier {
    /// Get expected latency in nanoseconds
    pub fn latency_ns(&self) -> u64 {
        match self {
            CxlTier::LocalDram => 100,   // ~100ns DRAM access
            CxlTier::CxlNear => 300,     // ~300ns CXL near
            CxlTier::CxlFar => 1000,     // ~1us CXL far
            CxlTier::Storage => 100_000, // ~100us SSD
        }
    }

    /// Get bandwidth in GB/s
    pub fn bandwidth_gbps(&self) -> u64 {
        match self {
            CxlTier::LocalDram => 200, // DDR5 ~200 GB/s per channel
            CxlTier::CxlNear => 64,    // CXL 2.0 ~64 GB/s
            CxlTier::CxlFar => 32,     // CXL far ~32 GB/s
            CxlTier::Storage => 7,     // NVMe ~7 GB/s
        }
    }

    /// Get cost per GB (relative scale)
    pub fn cost_per_gb(&self) -> u64 {
        match self {
            CxlTier::LocalDram => 100,
            CxlTier::CxlNear => 40,
            CxlTier::CxlFar => 20,
            CxlTier::Storage => 1,
        }
    }
}

/// Block temperature (access frequency)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockTemperature {
    /// Very hot (accessed very frequently)
    VeryHot,
    /// Hot (accessed frequently)
    Hot,
    /// Warm (accessed occasionally)
    Warm,
    /// Cold (rarely accessed)
    Cold,
    /// Frozen (not accessed in a long time)
    Frozen,
}

impl BlockTemperature {
    /// Calculate temperature from access count and time
    ///
    /// # Arguments
    /// * `access_count` - Number of accesses
    /// * `elapsed_seconds` - Time since first access
    pub fn from_access_pattern(access_count: u64, elapsed_seconds: u64) -> Self {
        if elapsed_seconds == 0 {
            return BlockTemperature::VeryHot;
        }

        let accesses_per_second = access_count / elapsed_seconds.max(1);

        match accesses_per_second {
            100.. => BlockTemperature::VeryHot,
            10..=99 => BlockTemperature::Hot,
            1..=9 => BlockTemperature::Warm,
            0 if access_count > 0 => BlockTemperature::Cold,
            _ => BlockTemperature::Frozen,
        }
    }

    /// Get recommended tier for this temperature
    pub fn recommended_tier(&self) -> CxlTier {
        match self {
            BlockTemperature::VeryHot => CxlTier::LocalDram,
            BlockTemperature::Hot => CxlTier::CxlNear,
            BlockTemperature::Warm => CxlTier::CxlFar,
            BlockTemperature::Cold | BlockTemperature::Frozen => CxlTier::Storage,
        }
    }
}

/// CXL block metadata
#[derive(Debug, Clone)]
pub struct CxlBlockMeta {
    /// Block identifier
    pub block_id: u64,
    /// Current tier
    pub current_tier: CxlTier,
    /// Access count
    pub access_count: u64,
    /// Last access timestamp
    pub last_access: u64,
    /// First access timestamp
    pub first_access: u64,
    /// Block temperature
    pub temperature: BlockTemperature,
}

impl CxlBlockMeta {
    /// Create new block metadata
    pub fn new(block_id: u64, tier: CxlTier, timestamp: u64) -> Self {
        Self {
            block_id,
            current_tier: tier,
            access_count: 0,
            last_access: timestamp,
            first_access: timestamp,
            temperature: BlockTemperature::Frozen,
        }
    }

    /// Update access statistics
    pub fn record_access(&mut self, timestamp: u64) {
        self.access_count += 1;
        self.last_access = timestamp;

        let elapsed = timestamp.saturating_sub(self.first_access);
        self.temperature = BlockTemperature::from_access_pattern(self.access_count, elapsed);
    }

    /// Check if block should be promoted to a higher tier
    pub fn should_promote(&self) -> bool {
        let recommended = self.temperature.recommended_tier();
        recommended < self.current_tier
    }

    /// Check if block should be demoted to a lower tier
    pub fn should_demote(&self, current_time: u64) -> bool {
        // Demote if not accessed in last 60 seconds
        const DEMOTE_THRESHOLD_SECONDS: u64 = 60;

        let time_since_access = current_time.saturating_sub(self.last_access);
        if time_since_access > DEMOTE_THRESHOLD_SECONDS {
            let recommended = self.temperature.recommended_tier();
            return recommended > self.current_tier;
        }

        false
    }
}

/// CXL tier statistics
#[derive(Debug, Clone, Default)]
pub struct CxlTierStats {
    /// Total capacity in bytes
    pub capacity: u64,
    /// Used bytes
    pub used: u64,
    /// Number of blocks
    pub block_count: usize,
    /// Promotion count
    pub promotions: u64,
    /// Demotion count
    pub demotions: u64,
}

lazy_static! {
    /// Global CXL memory manager
    static ref CXL_MANAGER: Mutex<CxlMemoryManager> = Mutex::new(CxlMemoryManager::new());
}

/// CXL memory tiering manager
pub struct CxlMemoryManager {
    /// Block metadata index
    blocks: BTreeMap<u64, CxlBlockMeta>,
    /// Tier statistics
    tier_stats: BTreeMap<CxlTier, CxlTierStats>,
    /// Current timestamp counter
    timestamp: u64,
    /// Total promotions
    total_promotions: u64,
    /// Total demotions
    total_demotions: u64,
}

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

impl CxlMemoryManager {
    /// Create new CXL memory manager
    pub fn new() -> Self {
        let mut tier_stats = BTreeMap::new();

        // Initialize tier capacities (example values)
        tier_stats.insert(
            CxlTier::LocalDram,
            CxlTierStats {
                capacity: 32 * 1024 * 1024 * 1024, // 32 GB
                ..Default::default()
            },
        );
        tier_stats.insert(
            CxlTier::CxlNear,
            CxlTierStats {
                capacity: 128 * 1024 * 1024 * 1024, // 128 GB
                ..Default::default()
            },
        );
        tier_stats.insert(
            CxlTier::CxlFar,
            CxlTierStats {
                capacity: 256 * 1024 * 1024 * 1024, // 256 GB
                ..Default::default()
            },
        );
        tier_stats.insert(
            CxlTier::Storage,
            CxlTierStats {
                capacity: 2 * 1024 * 1024 * 1024 * 1024, // 2 TB
                ..Default::default()
            },
        );

        Self {
            blocks: BTreeMap::new(),
            tier_stats,
            timestamp: 0,
            total_promotions: 0,
            total_demotions: 0,
        }
    }

    /// Initialize tier capacities
    ///
    /// # Arguments
    /// * `local_dram_gb` - Local DRAM capacity in GB
    /// * `cxl_near_gb` - CXL near capacity in GB
    /// * `cxl_far_gb` - CXL far capacity in GB
    /// * `storage_tb` - Storage capacity in TB
    pub fn init_tiers(
        &mut self,
        local_dram_gb: u64,
        cxl_near_gb: u64,
        cxl_far_gb: u64,
        storage_tb: u64,
    ) -> Result<(), &'static str> {
        let local_stats = self
            .tier_stats
            .get_mut(&CxlTier::LocalDram)
            .ok_or("LocalDram tier not initialized")?;
        local_stats.capacity = local_dram_gb * 1024 * 1024 * 1024;

        let near_stats = self
            .tier_stats
            .get_mut(&CxlTier::CxlNear)
            .ok_or("CxlNear tier not initialized")?;
        near_stats.capacity = cxl_near_gb * 1024 * 1024 * 1024;

        let far_stats = self
            .tier_stats
            .get_mut(&CxlTier::CxlFar)
            .ok_or("CxlFar tier not initialized")?;
        far_stats.capacity = cxl_far_gb * 1024 * 1024 * 1024;

        let storage_stats = self
            .tier_stats
            .get_mut(&CxlTier::Storage)
            .ok_or("Storage tier not initialized")?;
        storage_stats.capacity = storage_tb * 1024 * 1024 * 1024 * 1024;

        Ok(())
    }

    /// Allocate block in a tier
    ///
    /// # Arguments
    /// * `block_id` - Block identifier
    /// * `size` - Block size in bytes
    /// * `tier` - Target tier
    pub fn allocate_block(
        &mut self,
        block_id: u64,
        size: u64,
        tier: CxlTier,
    ) -> Result<(), &'static str> {
        let stats = self.tier_stats.get_mut(&tier).ok_or("Invalid tier")?;

        if stats.used + size > stats.capacity {
            return Err("Tier full");
        }

        self.timestamp += 1;
        let block = CxlBlockMeta::new(block_id, tier, self.timestamp);
        self.blocks.insert(block_id, block);

        stats.used += size;
        stats.block_count += 1;

        Ok(())
    }

    /// Record block access
    ///
    /// # Arguments
    /// * `block_id` - Block identifier
    pub fn access_block(&mut self, block_id: u64) {
        self.timestamp += 1;

        if let Some(block) = self.blocks.get_mut(&block_id) {
            block.record_access(self.timestamp);
        }
    }

    /// Promote block to higher tier
    ///
    /// # Arguments
    /// * `block_id` - Block identifier
    /// * `block_size` - Block size in bytes
    pub fn promote_block(
        &mut self,
        block_id: u64,
        block_size: u64,
    ) -> Result<CxlTier, &'static str> {
        let block = self.blocks.get_mut(&block_id).ok_or("Block not found")?;

        let new_tier = match block.current_tier {
            CxlTier::Storage => CxlTier::CxlFar,
            CxlTier::CxlFar => CxlTier::CxlNear,
            CxlTier::CxlNear => CxlTier::LocalDram,
            CxlTier::LocalDram => return Err("Already at highest tier"),
        };

        // Check if target tier has space
        let target_stats = self.tier_stats.get(&new_tier).ok_or("Invalid tier")?;
        if target_stats.used + block_size > target_stats.capacity {
            return Err("Target tier full");
        }

        // Update tier statistics
        let old_tier = block.current_tier;

        let old_stats = self
            .tier_stats
            .get_mut(&old_tier)
            .ok_or("Source tier not found")?;
        old_stats.used -= block_size;
        old_stats.block_count -= 1;
        old_stats.promotions += 1;

        let new_stats = self
            .tier_stats
            .get_mut(&new_tier)
            .ok_or("Target tier not found")?;
        new_stats.used += block_size;
        new_stats.block_count += 1;

        block.current_tier = new_tier;
        self.total_promotions += 1;

        Ok(new_tier)
    }

    /// Demote block to lower tier
    ///
    /// # Arguments
    /// * `block_id` - Block identifier
    /// * `block_size` - Block size in bytes
    pub fn demote_block(
        &mut self,
        block_id: u64,
        block_size: u64,
    ) -> Result<CxlTier, &'static str> {
        let block = self.blocks.get_mut(&block_id).ok_or("Block not found")?;

        let new_tier = match block.current_tier {
            CxlTier::LocalDram => CxlTier::CxlNear,
            CxlTier::CxlNear => CxlTier::CxlFar,
            CxlTier::CxlFar => CxlTier::Storage,
            CxlTier::Storage => return Err("Already at lowest tier"),
        };

        // Update tier statistics
        let old_tier = block.current_tier;

        let old_stats = self
            .tier_stats
            .get_mut(&old_tier)
            .ok_or("Source tier not found")?;
        old_stats.used -= block_size;
        old_stats.block_count -= 1;
        old_stats.demotions += 1;

        let new_stats = self
            .tier_stats
            .get_mut(&new_tier)
            .ok_or("Target tier not found")?;
        new_stats.used += block_size;
        new_stats.block_count += 1;

        block.current_tier = new_tier;
        self.total_demotions += 1;

        Ok(new_tier)
    }

    /// Run automatic tiering (promote/demote based on temperature)
    ///
    /// # Arguments
    /// * `block_size` - Block size in bytes
    ///
    /// # Returns
    /// (promotions, demotions)
    pub fn auto_tier(&mut self, block_size: u64) -> (u64, u64) {
        let mut promotions = 0;
        let mut demotions = 0;

        let block_ids: Vec<u64> = self.blocks.keys().copied().collect();

        for block_id in block_ids {
            // Get block or skip if it no longer exists
            let Some(block) = self.blocks.get(&block_id) else {
                continue;
            };

            let should_promote = block.should_promote();
            let should_demote = block.should_demote(self.timestamp);

            if should_promote {
                if self.promote_block(block_id, block_size).is_ok() {
                    promotions += 1;
                }
            } else if should_demote && self.demote_block(block_id, block_size).is_ok() {
                demotions += 1;
            }
        }

        (promotions, demotions)
    }

    /// Get tier statistics
    pub fn get_tier_stats(&self, tier: CxlTier) -> Option<CxlTierStats> {
        self.tier_stats.get(&tier).cloned()
    }

    /// Get global statistics
    ///
    /// # Returns
    /// (total_blocks, total_promotions, total_demotions)
    pub fn get_global_stats(&self) -> (usize, u64, u64) {
        (
            self.blocks.len(),
            self.total_promotions,
            self.total_demotions,
        )
    }
}

/// Global CXL operations
pub struct CxlEngine;

impl CxlEngine {
    /// Initialize CXL tiers
    pub fn init(
        local_dram_gb: u64,
        cxl_near_gb: u64,
        cxl_far_gb: u64,
        storage_tb: u64,
    ) -> Result<(), &'static str> {
        let mut mgr = CXL_MANAGER.lock();
        mgr.init_tiers(local_dram_gb, cxl_near_gb, cxl_far_gb, storage_tb)
    }

    /// Allocate block
    pub fn allocate(block_id: u64, size: u64, tier: CxlTier) -> Result<(), &'static str> {
        let mut mgr = CXL_MANAGER.lock();
        mgr.allocate_block(block_id, size, tier)
    }

    /// Access block
    pub fn access(block_id: u64) {
        let mut mgr = CXL_MANAGER.lock();
        mgr.access_block(block_id);
    }

    /// Promote block
    pub fn promote(block_id: u64, block_size: u64) -> Result<CxlTier, &'static str> {
        let mut mgr = CXL_MANAGER.lock();
        mgr.promote_block(block_id, block_size)
    }

    /// Demote block
    pub fn demote(block_id: u64, block_size: u64) -> Result<CxlTier, &'static str> {
        let mut mgr = CXL_MANAGER.lock();
        mgr.demote_block(block_id, block_size)
    }

    /// Run automatic tiering
    pub fn auto_tier(block_size: u64) -> (u64, u64) {
        let mut mgr = CXL_MANAGER.lock();
        mgr.auto_tier(block_size)
    }

    /// Get tier statistics
    pub fn tier_stats(tier: CxlTier) -> Option<CxlTierStats> {
        let mgr = CXL_MANAGER.lock();
        mgr.get_tier_stats(tier)
    }

    /// Get global statistics
    pub fn global_stats() -> (usize, u64, u64) {
        let mgr = CXL_MANAGER.lock();
        mgr.get_global_stats()
    }
}

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

    #[test]
    fn test_tier_properties() {
        assert_eq!(CxlTier::LocalDram.latency_ns(), 100);
        assert_eq!(CxlTier::CxlNear.latency_ns(), 300);
        assert!(CxlTier::LocalDram.bandwidth_gbps() > CxlTier::Storage.bandwidth_gbps());
    }

    #[test]
    fn test_block_temperature() {
        // Very hot: 100+ accesses per second
        let temp = BlockTemperature::from_access_pattern(1000, 5);
        assert_eq!(temp, BlockTemperature::VeryHot);

        // Hot: 10-99 accesses per second
        let temp = BlockTemperature::from_access_pattern(50, 5);
        assert_eq!(temp, BlockTemperature::Hot);

        // Warm: 1-9 accesses per second
        let temp = BlockTemperature::from_access_pattern(5, 5);
        assert_eq!(temp, BlockTemperature::Warm);

        // Cold: < 1 access per second
        let temp = BlockTemperature::from_access_pattern(1, 10);
        assert_eq!(temp, BlockTemperature::Cold);
    }

    #[test]
    fn test_recommended_tier() {
        assert_eq!(
            BlockTemperature::VeryHot.recommended_tier(),
            CxlTier::LocalDram
        );
        assert_eq!(BlockTemperature::Hot.recommended_tier(), CxlTier::CxlNear);
        assert_eq!(BlockTemperature::Warm.recommended_tier(), CxlTier::CxlFar);
        assert_eq!(BlockTemperature::Cold.recommended_tier(), CxlTier::Storage);
    }

    #[test]
    fn test_allocate_block() {
        let mut mgr = CxlMemoryManager::new();
        mgr.init_tiers(1, 4, 8, 1)
            .expect("test: operation should succeed");

        assert!(mgr.allocate_block(100, 4096, CxlTier::Storage).is_ok());
        assert_eq!(mgr.blocks.len(), 1);

        let stats = mgr
            .get_tier_stats(CxlTier::Storage)
            .expect("test: operation should succeed");
        assert_eq!(stats.used, 4096);
        assert_eq!(stats.block_count, 1);
    }

    #[test]
    fn test_access_tracking() {
        let mut mgr = CxlMemoryManager::new();
        mgr.allocate_block(100, 4096, CxlTier::Storage)
            .expect("test: operation should succeed");

        // Access 150 times - this will result in 150 accesses over 150 time units
        // 150 accesses / 150 time = 1 access/sec = Warm
        for _ in 0..150 {
            mgr.access_block(100);
        }

        let block = mgr
            .blocks
            .get(&100)
            .expect("test: operation should succeed");
        assert_eq!(block.access_count, 150);
        // With access_block() incrementing timestamp each time, we get 1 access/sec = Warm
        assert!(matches!(block.temperature, BlockTemperature::Warm));
    }

    #[test]
    fn test_promotion() {
        let mut mgr = CxlMemoryManager::new();
        mgr.init_tiers(1, 4, 8, 1)
            .expect("test: operation should succeed");
        mgr.allocate_block(100, 4096, CxlTier::Storage)
            .expect("test: operation should succeed");

        // Promote Storage -> CxlFar
        let new_tier = mgr
            .promote_block(100, 4096)
            .expect("test: operation should succeed");
        assert_eq!(new_tier, CxlTier::CxlFar);

        let block = mgr
            .blocks
            .get(&100)
            .expect("test: operation should succeed");
        assert_eq!(block.current_tier, CxlTier::CxlFar);
    }

    #[test]
    fn test_demotion() {
        let mut mgr = CxlMemoryManager::new();
        mgr.init_tiers(1, 4, 8, 1)
            .expect("test: operation should succeed");
        mgr.allocate_block(100, 4096, CxlTier::LocalDram)
            .expect("test: operation should succeed");

        // Demote LocalDram -> CxlNear
        let new_tier = mgr
            .demote_block(100, 4096)
            .expect("test: operation should succeed");
        assert_eq!(new_tier, CxlTier::CxlNear);

        let block = mgr
            .blocks
            .get(&100)
            .expect("test: operation should succeed");
        assert_eq!(block.current_tier, CxlTier::CxlNear);
    }

    #[test]
    fn test_auto_tiering() {
        let mut mgr = CxlMemoryManager::new();
        mgr.init_tiers(1, 4, 8, 1)
            .expect("test: operation should succeed");

        // Allocate cold block in DRAM
        mgr.allocate_block(100, 4096, CxlTier::LocalDram)
            .expect("test: operation should succeed");

        // Advance time without accessing (should demote)
        mgr.timestamp += 100;

        let (promotions, demotions) = mgr.auto_tier(4096);
        assert_eq!(promotions, 0);
        assert!(demotions > 0);
    }

    #[test]
    fn test_tier_full() {
        let mut mgr = CxlMemoryManager::new();
        mgr.init_tiers(0, 0, 0, 1)
            .expect("test: operation should succeed"); // Very small tiers

        // Fill LocalDram
        mgr.tier_stats
            .get_mut(&CxlTier::LocalDram)
            .expect("test: operation should succeed")
            .capacity = 4096;

        assert!(mgr.allocate_block(100, 4096, CxlTier::LocalDram).is_ok());
        assert!(mgr.allocate_block(101, 4096, CxlTier::LocalDram).is_err());
    }
}