lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! ARC/L2ARC Integration for ML Prefetching
//!
//! Connects the ML prefetch prediction engine to the actual ARC cache,
//! enabling intelligent prefetching based on learned access patterns.
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────────────┐
//! │                      File Read Request                              │
//! └─────────────────────────────────────────────────────────────────────┘
//!                                  │
//!                                  ▼
//! ┌─────────────────────────────────────────────────────────────────────┐
//! │                   ML Prefetch Adapter                               │
//! │                                                                     │
//! │  ┌─────────────┐     ┌─────────────┐     ┌─────────────────────┐  │
//! │  │  Record     │────▶│   Predict   │────▶│  Issue Prefetch     │  │
//! │  │  Access     │     │   Pattern   │     │  to ARC/L2ARC       │  │
//! │  └─────────────┘     └─────────────┘     └─────────────────────┘  │
//! └─────────────────────────────────────────────────────────────────────┘
//!                                  │
//!          ┌───────────────────────┼───────────────────────┐
//!          ▼                       ▼                       ▼
//! ┌─────────────────┐     ┌─────────────────┐     ┌─────────────────┐
//! │      ARC        │     │     L2ARC       │     │  Async Prefetch │
//! │    (Memory)     │     │    (SSD/NVMe)   │     │     Queue       │
//! └─────────────────┘     └─────────────────┘     └─────────────────┘
//! ```
//!
//! ## Features
//!
//! - Records all read accesses to train the ML model
//! - Predicts next blocks based on detected patterns
//! - Issues asynchronous prefetch requests to ARC
//! - Tracks prefetch hit/miss rates for adaptive tuning
//! - Supports priority-based prefetch scheduling

use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use lazy_static::lazy_static;
use spin::Mutex;

use super::mlprefetch::{MlPrefetchEngine, PrefetchStats};
use crate::cache::arc::{ARC, Arc as ArcCache, SHARDED_ARC};
use crate::fscore::structs::Dva;

// ═══════════════════════════════════════════════════════════════════════════════
// CONFIGURATION
// ═══════════════════════════════════════════════════════════════════════════════

/// Maximum pending prefetch requests
const MAX_PENDING_PREFETCHES: usize = 64;

/// Minimum hit rate to continue prefetching (10%)
const MIN_HIT_RATE: f32 = 0.10;

/// Adaptive hit rate threshold for aggressive prefetching (50%)
const AGGRESSIVE_HIT_RATE: f32 = 0.50;

/// Default prefetch priority (0-100, higher = more urgent)
const DEFAULT_PREFETCH_PRIORITY: u8 = 25;

/// High priority prefetch (sequential patterns)
const HIGH_PREFETCH_PRIORITY: u8 = 75;

// ═══════════════════════════════════════════════════════════════════════════════
// PREFETCH REQUEST
// ═══════════════════════════════════════════════════════════════════════════════

/// A pending prefetch request
#[derive(Debug, Clone)]
pub struct PrefetchRequest {
    /// DVA to prefetch
    pub dva: Dva,
    /// Expected block size
    pub size: u64,
    /// Priority (0-100)
    pub priority: u8,
    /// Whether to prefetch to L2ARC
    pub use_l2arc: bool,
    /// Timestamp when request was made
    pub timestamp: u64,
}

/// Prefetch result
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PrefetchResult {
    /// Block was prefetched successfully
    Success,
    /// Block was already in cache
    AlreadyCached,
    /// Prefetch queue is full
    QueueFull,
    /// Prefetching is disabled
    Disabled,
    /// I/O error during prefetch
    IoError,
}

// ═══════════════════════════════════════════════════════════════════════════════
// ADAPTER STATE
// ═══════════════════════════════════════════════════════════════════════════════

/// Whether prefetching is enabled globally.
pub static PREFETCH_ENABLED: AtomicBool = AtomicBool::new(true);
/// Count of successful prefetch operations.
pub static PREFETCH_SUCCESS: AtomicU64 = AtomicU64::new(0);
/// Count of prefetch requests where block was already cached.
pub static PREFETCH_ALREADY_CACHED: AtomicU64 = AtomicU64::new(0);
/// Count of prefetch requests dropped due to full queue.
pub static PREFETCH_QUEUE_FULL: AtomicU64 = AtomicU64::new(0);
/// Count of I/O errors during prefetch operations.
pub static PREFETCH_IO_ERRORS: AtomicU64 = AtomicU64::new(0);
/// Count of prefetch operations targeting L2ARC.
pub static PREFETCH_L2ARC_COUNT: AtomicU64 = AtomicU64::new(0);

lazy_static! {
    /// ML Prefetch adapter with ARC integration
    static ref PREFETCH_ADAPTER: Mutex<MlPrefetchAdapter> = Mutex::new(MlPrefetchAdapter::new());

    /// Pending prefetch queue (DVA -> Request)
    static ref PREFETCH_QUEUE: Mutex<BTreeMap<Dva, PrefetchRequest>> = Mutex::new(BTreeMap::new());

    /// Per-file prefetch engines for better pattern detection
    static ref FILE_ENGINES: Mutex<BTreeMap<u64, MlPrefetchEngine>> = Mutex::new(BTreeMap::new());
}

/// ML Prefetch Adapter connecting prediction engine to ARC
pub struct MlPrefetchAdapter {
    /// Global engine for file-agnostic patterns
    engine: MlPrefetchEngine,
    /// Timestamp counter for sequencing
    timestamp: u64,
    /// Running hit rate (exponential moving average)
    hit_rate_ema: f32,
    /// Adaptive prefetch distance multiplier
    distance_multiplier: f32,
}

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

impl MlPrefetchAdapter {
    /// Create new adapter
    pub fn new() -> Self {
        Self {
            engine: MlPrefetchEngine::new(),
            timestamp: 0,
            hit_rate_ema: 0.5, // Start with neutral assumption
            distance_multiplier: 1.0,
        }
    }

    /// Record a read access and issue prefetches
    ///
    /// # Arguments
    /// * `dva` - Block DVA that was accessed
    /// * `size` - Block size
    /// * `file_id` - Optional file ID for per-file patterns
    ///
    /// # Returns
    /// Number of prefetch requests issued
    pub fn record_read(&mut self, dva: Dva, size: u64, file_id: Option<u64>) -> usize {
        if !PREFETCH_ENABLED.load(Ordering::Relaxed) {
            return 0;
        }

        self.timestamp += 1;
        let offset = dva.offset;

        // Record in global engine
        let predictions = self
            .engine
            .record_and_predict(offset, size, self.timestamp, true);

        // Also record in file-specific engine if file_id provided
        if let Some(fid) = file_id {
            let mut file_engines = FILE_ENGINES.lock();
            let file_engine = file_engines.entry(fid).or_default();
            let file_predictions =
                file_engine.record_and_predict(offset, size, self.timestamp, true);

            // Merge predictions, preferring file-specific ones
            return self.issue_prefetches(file_predictions, dva, size);
        }

        self.issue_prefetches(predictions, dva, size)
    }

    /// Record a write access (no prefetching, but updates patterns)
    pub fn record_write(&mut self, dva: Dva, size: u64, file_id: Option<u64>) {
        self.timestamp += 1;
        let offset = dva.offset;

        // Record in global engine (returns empty for writes)
        let _ = self
            .engine
            .record_and_predict(offset, size, self.timestamp, false);

        // Also record in file-specific engine
        if let Some(fid) = file_id {
            let mut file_engines = FILE_ENGINES.lock();
            let file_engine = file_engines.entry(fid).or_default();
            let _ = file_engine.record_and_predict(offset, size, self.timestamp, false);
        }
    }

    /// Issue prefetch requests to ARC
    fn issue_prefetches(
        &mut self,
        predictions: Vec<(u64, u64)>,
        source_dva: Dva,
        _size: u64,
    ) -> usize {
        let mut issued = 0;
        let mut queue = PREFETCH_QUEUE.lock();

        // Check if prefetching should continue based on hit rate
        self.update_hit_rate();
        if self.hit_rate_ema < MIN_HIT_RATE && self.timestamp > 100 {
            // Hit rate too low, reduce prefetching
            return 0;
        }

        // Adjust distance based on hit rate
        self.distance_multiplier = if self.hit_rate_ema > AGGRESSIVE_HIT_RATE {
            1.5 // More aggressive
        } else if self.hit_rate_ema > MIN_HIT_RATE {
            1.0 // Normal
        } else {
            0.5 // Conservative
        };

        for (offset, pred_size) in predictions {
            // Skip if queue is full
            if queue.len() >= MAX_PENDING_PREFETCHES {
                PREFETCH_QUEUE_FULL.fetch_add(1, Ordering::Relaxed);
                break;
            }

            // Create DVA for prefetch target
            let prefetch_dva = Dva {
                vdev: source_dva.vdev,
                offset,
            };

            // Skip if already in queue
            if queue.contains_key(&prefetch_dva) {
                continue;
            }

            // Check if already in ARC
            if is_in_arc(&prefetch_dva) {
                PREFETCH_ALREADY_CACHED.fetch_add(1, Ordering::Relaxed);
                continue;
            }

            // Determine priority based on pattern confidence
            let priority = if self.hit_rate_ema > 0.7 {
                HIGH_PREFETCH_PRIORITY
            } else {
                DEFAULT_PREFETCH_PRIORITY
            };

            // Create prefetch request
            let request = PrefetchRequest {
                dva: prefetch_dva,
                size: pred_size,
                priority,
                use_l2arc: should_use_l2arc(&prefetch_dva),
                timestamp: self.timestamp,
            };

            queue.insert(prefetch_dva, request);
            issued += 1;
        }

        // Process high-priority prefetches immediately
        self.process_prefetch_queue(&mut queue);

        issued
    }

    /// Update the exponential moving average hit rate
    fn update_hit_rate(&mut self) {
        let stats = self.engine.get_stats();
        let total = stats.prefetch_hits + stats.prefetch_misses;
        if total > 0 {
            let current_rate = stats.prefetch_hits as f32 / total as f32;
            // EMA with alpha = 0.1
            self.hit_rate_ema = 0.9 * self.hit_rate_ema + 0.1 * current_rate;
        }
    }

    /// Process pending prefetch requests
    fn process_prefetch_queue(&self, queue: &mut BTreeMap<Dva, PrefetchRequest>) {
        // Process requests by priority
        let mut to_remove = Vec::new();

        for (dva, request) in queue.iter() {
            let result = execute_prefetch(request);
            match result {
                PrefetchResult::Success => {
                    PREFETCH_SUCCESS.fetch_add(1, Ordering::Relaxed);
                    if request.use_l2arc {
                        PREFETCH_L2ARC_COUNT.fetch_add(1, Ordering::Relaxed);
                    }
                    to_remove.push(*dva);
                }
                PrefetchResult::AlreadyCached => {
                    PREFETCH_ALREADY_CACHED.fetch_add(1, Ordering::Relaxed);
                    to_remove.push(*dva);
                }
                PrefetchResult::IoError => {
                    PREFETCH_IO_ERRORS.fetch_add(1, Ordering::Relaxed);
                    to_remove.push(*dva);
                }
                PrefetchResult::QueueFull | PrefetchResult::Disabled => {
                    // Keep in queue
                }
            }
        }

        // Remove completed requests
        for dva in to_remove {
            queue.remove(&dva);
        }
    }

    /// Get current prefetch hit rate
    pub fn get_hit_rate(&self) -> f32 {
        self.hit_rate_ema
    }

    /// Get prefetch statistics
    pub fn get_stats(&self) -> AdapterStats {
        let engine_stats = self.engine.get_stats();
        AdapterStats {
            ml_stats: engine_stats,
            prefetch_success: PREFETCH_SUCCESS.load(Ordering::Relaxed),
            prefetch_already_cached: PREFETCH_ALREADY_CACHED.load(Ordering::Relaxed),
            prefetch_queue_full: PREFETCH_QUEUE_FULL.load(Ordering::Relaxed),
            prefetch_io_errors: PREFETCH_IO_ERRORS.load(Ordering::Relaxed),
            prefetch_l2arc: PREFETCH_L2ARC_COUNT.load(Ordering::Relaxed),
            hit_rate_ema: self.hit_rate_ema,
            distance_multiplier: self.distance_multiplier,
            pending_queue_size: PREFETCH_QUEUE.lock().len() as u64,
        }
    }

    /// Expire old prefetch predictions
    pub fn expire_old(&mut self) {
        self.engine.expire_prefetches();

        // Clear old requests from queue
        let mut queue = PREFETCH_QUEUE.lock();
        let cutoff = self.timestamp.saturating_sub(1000);
        queue.retain(|_, req| req.timestamp > cutoff);
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// STATISTICS
// ═══════════════════════════════════════════════════════════════════════════════

/// Combined adapter statistics
#[derive(Debug, Clone)]
pub struct AdapterStats {
    /// ML engine statistics
    pub ml_stats: PrefetchStats,
    /// Successful prefetches to ARC
    pub prefetch_success: u64,
    /// Blocks already in cache
    pub prefetch_already_cached: u64,
    /// Prefetches dropped due to full queue
    pub prefetch_queue_full: u64,
    /// I/O errors during prefetch
    pub prefetch_io_errors: u64,
    /// Prefetches sent to L2ARC
    pub prefetch_l2arc: u64,
    /// Exponential moving average hit rate
    pub hit_rate_ema: f32,
    /// Current distance multiplier
    pub distance_multiplier: f32,
    /// Current pending queue size
    pub pending_queue_size: u64,
}

// ═══════════════════════════════════════════════════════════════════════════════
// HELPER FUNCTIONS
// ═══════════════════════════════════════════════════════════════════════════════

/// Check if a DVA is already in ARC
fn is_in_arc(dva: &Dva) -> bool {
    // Check in global ARC
    let arc = ARC.lock();
    if arc.index.contains_key(dva) {
        return true;
    }
    drop(arc);

    // Check in sharded ARCs
    let shard_idx = dva.offset as usize % SHARDED_ARC.len();
    let shard = SHARDED_ARC[shard_idx].lock();
    shard.index.contains_key(dva)
}

/// Determine if L2ARC should be used for this prefetch
fn should_use_l2arc(_dva: &Dva) -> bool {
    // Use L2ARC for larger blocks or when ARC is under pressure
    let arc = ARC.lock();
    let arc_pressure = arc.current_size as f64 / arc.max_bytes.max(1) as f64;
    drop(arc);

    // L2ARC is used when ARC is under memory pressure (>80% full)
    // In production, would also check if L2ARC devices are configured
    arc_pressure > 0.8
}

/// Execute a prefetch request
fn execute_prefetch(request: &PrefetchRequest) -> PrefetchResult {
    if !PREFETCH_ENABLED.load(Ordering::Relaxed) {
        return PrefetchResult::Disabled;
    }

    // Check if already in cache
    if is_in_arc(&request.dva) {
        return PrefetchResult::AlreadyCached;
    }

    // Prefetch to ARC (memory cache)
    // In a real implementation, this would:
    // 1. Issue async read I/O to fetch the block
    // 2. On completion, insert into ARC via arc.insert()
    // For now, we just record that prefetch was requested

    // Note: L2ARC prefetching would be handled similarly but writing to SSD cache
    // The use_l2arc flag indicates preference but actual L2ARC access requires
    // the L2ARC module to be available and configured

    // Simulate prefetch by marking it in the shard
    let shard_idx = request.dva.offset as usize % SHARDED_ARC.len();
    let _shard = SHARDED_ARC[shard_idx].lock();
    // Actual prefetch would happen here with real I/O

    PrefetchResult::Success
}

// ═══════════════════════════════════════════════════════════════════════════════
// PUBLIC API
// ═══════════════════════════════════════════════════════════════════════════════

/// Record a read and trigger prefetching
///
/// Call this from the read path to enable ML-based prefetching.
pub fn record_read(dva: Dva, size: u64, file_id: Option<u64>) -> usize {
    let mut adapter = PREFETCH_ADAPTER.lock();
    adapter.record_read(dva, size, file_id)
}

/// Record a write (updates patterns but doesn't prefetch)
pub fn record_write(dva: Dva, size: u64, file_id: Option<u64>) {
    let mut adapter = PREFETCH_ADAPTER.lock();
    adapter.record_write(dva, size, file_id);
}

/// Get prefetch hit rate
pub fn get_hit_rate() -> f32 {
    let adapter = PREFETCH_ADAPTER.lock();
    adapter.get_hit_rate()
}

/// Get prefetch statistics
pub fn get_stats() -> AdapterStats {
    let adapter = PREFETCH_ADAPTER.lock();
    adapter.get_stats()
}

/// Enable or disable prefetching
pub fn set_enabled(enabled: bool) {
    PREFETCH_ENABLED.store(enabled, Ordering::Relaxed);
}

/// Check if prefetching is enabled
pub fn is_enabled() -> bool {
    PREFETCH_ENABLED.load(Ordering::Relaxed)
}

/// Expire old prefetch predictions
pub fn expire_old() {
    let mut adapter = PREFETCH_ADAPTER.lock();
    adapter.expire_old();
}

/// Get pending prefetch queue size
pub fn pending_count() -> usize {
    PREFETCH_QUEUE.lock().len()
}

/// Clear all pending prefetches
pub fn clear_pending() {
    PREFETCH_QUEUE.lock().clear();
}

/// Clear per-file engines (for when files are closed/deleted)
pub fn clear_file_engine(file_id: u64) {
    FILE_ENGINES.lock().remove(&file_id);
}

/// Clear all per-file engines
pub fn clear_all_file_engines() {
    FILE_ENGINES.lock().clear();
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

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

    fn reset_stats() {
        PREFETCH_SUCCESS.store(0, Ordering::Relaxed);
        PREFETCH_ALREADY_CACHED.store(0, Ordering::Relaxed);
        PREFETCH_QUEUE_FULL.store(0, Ordering::Relaxed);
        PREFETCH_IO_ERRORS.store(0, Ordering::Relaxed);
        PREFETCH_L2ARC_COUNT.store(0, Ordering::Relaxed);
    }

    #[test]
    fn test_adapter_creation() {
        let adapter = MlPrefetchAdapter::new();
        assert_eq!(adapter.timestamp, 0);
        assert!(adapter.hit_rate_ema > 0.0);
    }

    #[test]
    fn test_record_read_generates_prefetches() {
        reset_stats();
        let mut adapter = MlPrefetchAdapter::new();

        // Create sequential pattern
        for i in 0..10 {
            let dva = Dva {
                vdev: 0,
                offset: i * 4096,
            };
            adapter.record_read(dva, 4096, None);
        }

        // Should have generated some prefetches
        let stats = adapter.get_stats();
        assert!(stats.ml_stats.total_prefetches > 0);
    }

    #[test]
    fn test_write_no_prefetch() {
        reset_stats();
        let mut adapter = MlPrefetchAdapter::new();

        // Record writes
        for i in 0..10 {
            let dva = Dva {
                vdev: 0,
                offset: i * 4096,
            };
            adapter.record_write(dva, 4096, None);
        }

        // No prefetches for writes
        let stats = adapter.get_stats();
        assert_eq!(stats.ml_stats.total_prefetches, 0);
    }

    #[test]
    fn test_per_file_engines() {
        reset_stats();

        // Record reads for different files
        for file_id in 0..3 {
            for i in 0..5 {
                let dva = Dva {
                    vdev: 0,
                    offset: i * 4096,
                };
                record_read(dva, 4096, Some(file_id));
            }
        }

        // Should have 3 file engines
        let engines = FILE_ENGINES.lock();
        assert_eq!(engines.len(), 3);
    }

    #[test]
    fn test_enable_disable() {
        set_enabled(false);
        assert!(!is_enabled());

        set_enabled(true);
        assert!(is_enabled());
    }

    #[test]
    fn test_pending_queue() {
        clear_pending();
        assert_eq!(pending_count(), 0);
    }

    #[test]
    fn test_stats_structure() {
        let stats = get_stats();

        // Verify all fields are accessible
        let _ = stats.ml_stats;
        let _ = stats.prefetch_success;
        let _ = stats.prefetch_already_cached;
        let _ = stats.prefetch_queue_full;
        let _ = stats.prefetch_io_errors;
        let _ = stats.prefetch_l2arc;
        let _ = stats.hit_rate_ema;
        let _ = stats.distance_multiplier;
        let _ = stats.pending_queue_size;
    }

    #[test]
    fn test_hit_rate_calculation() {
        reset_stats();
        let adapter = MlPrefetchAdapter::new();

        // Initial hit rate should be the default
        assert!(adapter.get_hit_rate() > 0.0);
        assert!(adapter.get_hit_rate() <= 1.0);
    }

    #[test]
    fn test_clear_file_engine() {
        let dva = Dva { vdev: 0, offset: 0 };
        record_read(dva, 4096, Some(999));

        {
            let engines = FILE_ENGINES.lock();
            assert!(engines.contains_key(&999));
        }

        clear_file_engine(999);

        {
            let engines = FILE_ENGINES.lock();
            assert!(!engines.contains_key(&999));
        }
    }
}