lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// QoS & Throttling
// Per-dataset bandwidth limits and I/O priority management.

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

/// I/O priority levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum IoPriority {
    /// Critical (highest priority) - system-critical operations
    Critical = 0,
    /// High - interactive workloads
    High = 1,
    /// Normal - default priority
    Normal = 2,
    /// Low - background tasks
    Low = 3,
    /// Idle - lowest priority, only when system is idle
    Idle = 4,
}

impl IoPriority {
    /// Get weight for scheduling (higher = more CPU time)
    pub fn weight(&self) -> u32 {
        match self {
            IoPriority::Critical => 1000,
            IoPriority::High => 500,
            IoPriority::Normal => 100,
            IoPriority::Low => 50,
            IoPriority::Idle => 10,
        }
    }

    /// Get name
    pub fn name(&self) -> &'static str {
        match self {
            IoPriority::Critical => "critical",
            IoPriority::High => "high",
            IoPriority::Normal => "normal",
            IoPriority::Low => "low",
            IoPriority::Idle => "idle",
        }
    }
}

/// Token bucket for bandwidth throttling
#[derive(Debug, Clone)]
pub struct TokenBucket {
    /// Maximum tokens (burst size)
    capacity: u64,
    /// Current tokens available
    tokens: u64,
    /// Refill rate (bytes per second)
    refill_rate: u64,
    /// Last refill timestamp
    last_refill: u64,
}

impl TokenBucket {
    /// Create new token bucket
    ///
    /// # Arguments
    /// * `capacity` - Burst size in bytes
    /// * `refill_rate` - Sustained rate in bytes/sec
    pub fn new(capacity: u64, refill_rate: u64) -> Self {
        Self {
            capacity,
            tokens: capacity, // Start full
            refill_rate,
            last_refill: 0,
        }
    }

    /// Refill tokens based on elapsed time
    ///
    /// # Arguments
    /// * `current_time` - Current timestamp in seconds
    pub fn refill(&mut self, current_time: u64) {
        let elapsed = current_time.saturating_sub(self.last_refill);
        if elapsed > 0 {
            let new_tokens = elapsed * self.refill_rate;
            self.tokens = (self.tokens + new_tokens).min(self.capacity);
            self.last_refill = current_time;
        }
    }

    /// Try to consume tokens for an I/O operation
    ///
    /// # Arguments
    /// * `bytes` - Number of bytes to consume
    /// * `current_time` - Current timestamp
    ///
    /// # Returns
    /// * `Ok(())` if tokens available
    /// * `Err(wait_time)` with seconds to wait if not enough tokens
    pub fn consume(&mut self, bytes: u64, current_time: u64) -> Result<(), u64> {
        self.refill(current_time);

        if self.tokens >= bytes {
            self.tokens -= bytes;
            Ok(())
        } else {
            // Calculate wait time needed
            let needed = bytes - self.tokens;
            let wait_time = needed.div_ceil(self.refill_rate); // Ceiling division
            Err(wait_time)
        }
    }

    /// Get current fill level (0.0 - 1.0)
    pub fn fill_level(&self) -> f32 {
        if self.capacity == 0 {
            return 0.0;
        }
        self.tokens as f32 / self.capacity as f32
    }
}

/// QoS policy for a dataset
#[derive(Debug, Clone)]
pub struct QosPolicy {
    /// Dataset ID
    pub dataset_id: u64,
    /// I/O priority
    pub priority: IoPriority,
    /// Read bandwidth limit (bytes/sec, 0 = unlimited)
    pub read_limit_bps: u64,
    /// Write bandwidth limit (bytes/sec, 0 = unlimited)
    pub write_limit_bps: u64,
    /// Read token bucket
    read_bucket: Option<TokenBucket>,
    /// Write token bucket
    write_bucket: Option<TokenBucket>,
}

impl QosPolicy {
    /// Create new QoS policy
    pub fn new(dataset_id: u64, priority: IoPriority) -> Self {
        Self {
            dataset_id,
            priority,
            read_limit_bps: 0,  // Unlimited by default
            write_limit_bps: 0, // Unlimited by default
            read_bucket: None,
            write_bucket: None,
        }
    }

    /// Set read bandwidth limit
    ///
    /// # Arguments
    /// * `bytes_per_sec` - Bandwidth limit (0 = unlimited)
    /// * `burst_size` - Burst size in bytes (default: 10x rate)
    pub fn set_read_limit(&mut self, bytes_per_sec: u64, burst_size: Option<u64>) {
        self.read_limit_bps = bytes_per_sec;
        if bytes_per_sec > 0 {
            let burst = burst_size.unwrap_or(bytes_per_sec * 10);
            self.read_bucket = Some(TokenBucket::new(burst, bytes_per_sec));
        } else {
            self.read_bucket = None;
        }
    }

    /// Set write bandwidth limit
    pub fn set_write_limit(&mut self, bytes_per_sec: u64, burst_size: Option<u64>) {
        self.write_limit_bps = bytes_per_sec;
        if bytes_per_sec > 0 {
            let burst = burst_size.unwrap_or(bytes_per_sec * 10);
            self.write_bucket = Some(TokenBucket::new(burst, bytes_per_sec));
        } else {
            self.write_bucket = None;
        }
    }

    /// Check if read is allowed (throttle if needed)
    ///
    /// # Arguments
    /// * `bytes` - Number of bytes to read
    /// * `current_time` - Current timestamp
    ///
    /// # Returns
    /// * `Ok(())` if allowed
    /// * `Err(wait_time)` if throttled
    pub fn check_read(&mut self, bytes: u64, current_time: u64) -> Result<(), u64> {
        if let Some(bucket) = &mut self.read_bucket {
            bucket.consume(bytes, current_time)
        } else {
            Ok(()) // No limit
        }
    }

    /// Check if write is allowed
    pub fn check_write(&mut self, bytes: u64, current_time: u64) -> Result<(), u64> {
        if let Some(bucket) = &mut self.write_bucket {
            bucket.consume(bytes, current_time)
        } else {
            Ok(()) // No limit
        }
    }
}

/// I/O statistics for QoS monitoring
#[derive(Debug, Clone, Default)]
pub struct QosStats {
    /// Total bytes read
    pub bytes_read: u64,
    /// Total bytes written
    pub bytes_written: u64,
    /// Number of times throttled
    pub throttle_count: u64,
    /// Total throttle time (seconds)
    pub throttle_time: u64,
    /// I/O operations completed
    pub ops_completed: u64,
}

lazy_static! {
    /// Global QoS manager
    static ref QOS_MANAGER: Mutex<QosManager> = Mutex::new(QosManager::new());
}

/// QoS manager
pub struct QosManager {
    /// QoS policies per dataset
    policies: BTreeMap<u64, QosPolicy>,
    /// Statistics per dataset
    stats: BTreeMap<u64, QosStats>,
    /// Default priority for datasets without policy
    default_priority: IoPriority,
}

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

impl QosManager {
    /// Create new QoS manager
    pub fn new() -> Self {
        Self {
            policies: BTreeMap::new(),
            stats: BTreeMap::new(),
            default_priority: IoPriority::Normal,
        }
    }

    /// Set QoS policy for dataset
    pub fn set_policy(&mut self, policy: QosPolicy) {
        self.policies.insert(policy.dataset_id, policy);
    }

    /// Get QoS policy for dataset
    pub fn get_policy(&self, dataset_id: u64) -> Option<&QosPolicy> {
        self.policies.get(&dataset_id)
    }

    /// Get mutable QoS policy
    pub fn get_policy_mut(&mut self, dataset_id: u64) -> Option<&mut QosPolicy> {
        self.policies.get_mut(&dataset_id)
    }

    /// Set priority for dataset
    pub fn set_priority(&mut self, dataset_id: u64, priority: IoPriority) {
        self.policies
            .entry(dataset_id)
            .or_insert_with(|| QosPolicy::new(dataset_id, priority))
            .priority = priority;
    }

    /// Set read bandwidth limit
    pub fn set_read_limit(&mut self, dataset_id: u64, bytes_per_sec: u64) {
        let policy = self
            .policies
            .entry(dataset_id)
            .or_insert_with(|| QosPolicy::new(dataset_id, self.default_priority));
        policy.set_read_limit(bytes_per_sec, None);
    }

    /// Set write bandwidth limit
    pub fn set_write_limit(&mut self, dataset_id: u64, bytes_per_sec: u64) {
        let policy = self
            .policies
            .entry(dataset_id)
            .or_insert_with(|| QosPolicy::new(dataset_id, self.default_priority));
        policy.set_write_limit(bytes_per_sec, None);
    }

    /// Check and throttle read operation
    ///
    /// # Returns
    /// * `Ok(())` if allowed
    /// * `Err(wait_time)` if should throttle
    pub fn throttle_read(
        &mut self,
        dataset_id: u64,
        bytes: u64,
        current_time: u64,
    ) -> Result<(), u64> {
        if let Some(policy) = self.policies.get_mut(&dataset_id) {
            let result = policy.check_read(bytes, current_time);

            // Update stats
            let stats = self.stats.entry(dataset_id).or_default();
            stats.bytes_read += bytes;

            if let Err(wait) = result {
                stats.throttle_count += 1;
                stats.throttle_time += wait;
            }

            result
        } else {
            // No policy = no throttling
            let stats = self.stats.entry(dataset_id).or_default();
            stats.bytes_read += bytes;
            Ok(())
        }
    }

    /// Check and throttle write operation
    pub fn throttle_write(
        &mut self,
        dataset_id: u64,
        bytes: u64,
        current_time: u64,
    ) -> Result<(), u64> {
        if let Some(policy) = self.policies.get_mut(&dataset_id) {
            let result = policy.check_write(bytes, current_time);

            // Update stats
            let stats = self.stats.entry(dataset_id).or_default();
            stats.bytes_written += bytes;

            if let Err(wait) = result {
                stats.throttle_count += 1;
                stats.throttle_time += wait;
            }

            result
        } else {
            let stats = self.stats.entry(dataset_id).or_default();
            stats.bytes_written += bytes;
            Ok(())
        }
    }

    /// Record I/O completion
    pub fn record_completion(&mut self, dataset_id: u64) {
        self.stats.entry(dataset_id).or_default().ops_completed += 1;
    }

    /// Get statistics for dataset
    pub fn get_stats(&self, dataset_id: u64) -> Option<QosStats> {
        self.stats.get(&dataset_id).cloned()
    }

    /// Get priority for dataset
    pub fn get_priority(&self, dataset_id: u64) -> IoPriority {
        self.policies
            .get(&dataset_id)
            .map(|p| p.priority)
            .unwrap_or(self.default_priority)
    }
}

/// Global QoS operations
pub struct QosEngine;

impl QosEngine {
    /// Set I/O priority for dataset
    pub fn set_priority(dataset_id: u64, priority: IoPriority) {
        let mut mgr = QOS_MANAGER.lock();
        mgr.set_priority(dataset_id, priority);
    }

    /// Set read bandwidth limit
    pub fn set_read_limit(dataset_id: u64, bytes_per_sec: u64) {
        let mut mgr = QOS_MANAGER.lock();
        mgr.set_read_limit(dataset_id, bytes_per_sec);
    }

    /// Set write bandwidth limit
    pub fn set_write_limit(dataset_id: u64, bytes_per_sec: u64) {
        let mut mgr = QOS_MANAGER.lock();
        mgr.set_write_limit(dataset_id, bytes_per_sec);
    }

    /// Throttle read operation
    pub fn throttle_read(dataset_id: u64, bytes: u64, current_time: u64) -> Result<(), u64> {
        let mut mgr = QOS_MANAGER.lock();
        mgr.throttle_read(dataset_id, bytes, current_time)
    }

    /// Throttle write operation
    pub fn throttle_write(dataset_id: u64, bytes: u64, current_time: u64) -> Result<(), u64> {
        let mut mgr = QOS_MANAGER.lock();
        mgr.throttle_write(dataset_id, bytes, current_time)
    }

    /// Record operation completion
    pub fn record_completion(dataset_id: u64) {
        let mut mgr = QOS_MANAGER.lock();
        mgr.record_completion(dataset_id);
    }

    /// Get statistics
    pub fn stats(dataset_id: u64) -> Option<QosStats> {
        let mgr = QOS_MANAGER.lock();
        mgr.get_stats(dataset_id)
    }

    /// Get priority
    pub fn get_priority(dataset_id: u64) -> IoPriority {
        let mgr = QOS_MANAGER.lock();
        mgr.get_priority(dataset_id)
    }
}

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

    #[test]
    fn test_priority_weights() {
        assert!(IoPriority::Critical.weight() > IoPriority::High.weight());
        assert!(IoPriority::High.weight() > IoPriority::Normal.weight());
        assert!(IoPriority::Normal.weight() > IoPriority::Low.weight());
        assert!(IoPriority::Low.weight() > IoPriority::Idle.weight());
    }

    #[test]
    fn test_token_bucket_basic() {
        let mut bucket = TokenBucket::new(1000, 100); // 1KB capacity, 100 B/s refill

        // Should allow 1000 bytes initially
        assert!(bucket.consume(1000, 0).is_ok());

        // Should be empty now
        assert!(bucket.consume(1, 0).is_err());
    }

    #[test]
    fn test_token_bucket_refill() {
        let mut bucket = TokenBucket::new(1000, 100);

        // Consume all tokens
        bucket
            .consume(1000, 0)
            .expect("test: operation should succeed");

        // After 5 seconds, should have 500 tokens
        assert!(bucket.consume(500, 5).is_ok());

        // Should be empty again
        assert!(bucket.consume(1, 5).is_err());
    }

    #[test]
    fn test_token_bucket_burst() {
        let mut bucket = TokenBucket::new(10000, 1000); // 10KB burst, 1KB/s sustained

        // Can burst up to 10KB
        assert!(bucket.consume(10000, 0).is_ok());

        // But sustained rate is only 1KB/s
        assert!(bucket.consume(5000, 3).is_err()); // Only 3KB refilled
    }

    #[test]
    fn test_qos_policy_unlimited() {
        let mut policy = QosPolicy::new(1, IoPriority::Normal);

        // No limits set = unlimited
        assert!(policy.check_read(1000000, 0).is_ok());
        assert!(policy.check_write(1000000, 0).is_ok());
    }

    #[test]
    fn test_qos_policy_limits() {
        let mut policy = QosPolicy::new(1, IoPriority::Normal);
        policy.set_read_limit(1000, None); // 1KB/s

        // Should allow burst
        assert!(policy.check_read(10000, 0).is_ok());

        // Should throttle now
        assert!(policy.check_read(1, 0).is_err());
    }

    #[test]
    fn test_qos_manager_priority() {
        let mut mgr = QosManager::new();

        mgr.set_priority(100, IoPriority::High);
        assert_eq!(mgr.get_priority(100), IoPriority::High);

        // Default for unset dataset
        assert_eq!(mgr.get_priority(200), IoPriority::Normal);
    }

    #[test]
    fn test_qos_manager_throttling() {
        let mut mgr = QosManager::new();
        mgr.set_read_limit(100, 1000); // 1KB/s for dataset 100 (burst = 10KB)

        // First read should work (burst available) - consumes all tokens
        assert!(mgr.throttle_read(100, 10000, 0).is_ok());

        // Next read should throttle (no tokens left)
        assert!(mgr.throttle_read(100, 1000, 0).is_err());

        // After time passes, should work again
        assert!(mgr.throttle_read(100, 1000, 2).is_ok());
    }

    #[test]
    fn test_qos_statistics() {
        let mut mgr = QosManager::new();
        mgr.set_read_limit(100, 1000); // 1KB/s (burst = 10KB)

        // Perform some I/O
        let _ = mgr.throttle_read(100, 10000, 0); // Use full burst
        let _ = mgr.throttle_read(100, 1000, 0); // This will throttle
        mgr.record_completion(100);

        let stats = mgr.get_stats(100).expect("test: operation should succeed");
        assert_eq!(stats.bytes_read, 11000);
        assert_eq!(stats.throttle_count, 1);
        assert_eq!(stats.ops_completed, 1);
    }

    #[test]
    fn test_fill_level() {
        let mut bucket = TokenBucket::new(1000, 100);

        assert!((bucket.fill_level() - 1.0).abs() < 0.01); // Full

        bucket
            .consume(500, 0)
            .expect("test: operation should succeed");
        assert!((bucket.fill_level() - 0.5).abs() < 0.01); // Half
    }
}