lcpfs 2026.1.102

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

//! Quota storage and management.
//!
//! This module provides storage for quota configuration and usage tracking.

use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use lazy_static::lazy_static;
use spin::Mutex;

use super::types::{
    Quota, QuotaCheckResult, QuotaError, QuotaKey, QuotaLimits, QuotaReport, QuotaReportEntry,
    QuotaResult, QuotaStatus, QuotaType, QuotaUsage,
};

// ═══════════════════════════════════════════════════════════════════════════════
// GLOBAL QUOTA STORE
// ═══════════════════════════════════════════════════════════════════════════════

lazy_static! {
    /// Global quota storage.
    static ref QUOTA_STORE: Mutex<QuotaStore> = Mutex::new(QuotaStore::new());
}

/// Quota store for all datasets.
#[derive(Debug, Default)]
struct QuotaStore {
    /// Quotas by dataset.
    datasets: BTreeMap<String, DatasetQuotas>,
}

impl QuotaStore {
    /// Create a new quota store.
    fn new() -> Self {
        Self::default()
    }

    /// Get or create dataset quotas.
    fn get_or_create(&mut self, dataset: &str) -> &mut DatasetQuotas {
        if !self.datasets.contains_key(dataset) {
            self.datasets
                .insert(dataset.to_string(), DatasetQuotas::new());
        }
        self.datasets.get_mut(dataset).unwrap()
    }

    /// Get dataset quotas.
    fn get(&self, dataset: &str) -> Option<&DatasetQuotas> {
        self.datasets.get(dataset)
    }

    /// Get dataset quotas mutably.
    fn get_mut(&mut self, dataset: &str) -> Option<&mut DatasetQuotas> {
        self.datasets.get_mut(dataset)
    }
}

/// Quotas for a single dataset.
#[derive(Debug, Default)]
struct DatasetQuotas {
    /// Quota entries by key.
    quotas: BTreeMap<QuotaKey, Quota>,
    /// Usage tracking (separate from quotas for quick updates).
    usage: BTreeMap<QuotaKey, QuotaUsage>,
}

impl DatasetQuotas {
    /// Create new dataset quotas.
    fn new() -> Self {
        Self::default()
    }
}

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

/// Set quota limits for a user/group.
pub fn set_quota(dataset: &str, key: QuotaKey, limits: QuotaLimits) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    let ds = store.get_or_create(dataset);

    if let Some(quota) = ds.quotas.get_mut(&key) {
        quota.limits = limits;
    } else {
        ds.quotas.insert(key, Quota::new(key, limits));
    }

    Ok(())
}

/// Get quota for a user/group.
pub fn get_quota(dataset: &str, key: QuotaKey) -> QuotaResult<Quota> {
    let store = QUOTA_STORE.lock();
    let ds = store
        .get(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    ds.quotas
        .get(&key)
        .cloned()
        .ok_or(QuotaError::NotFound(key))
}

/// Remove quota for a user/group.
pub fn remove_quota(dataset: &str, key: QuotaKey) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    let ds = store
        .get_mut(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    ds.quotas.remove(&key);
    ds.usage.remove(&key);

    Ok(())
}

/// Get current usage for a user/group.
pub fn get_usage(dataset: &str, key: QuotaKey) -> QuotaResult<QuotaUsage> {
    let store = QUOTA_STORE.lock();
    let ds = store
        .get(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    Ok(ds.usage.get(&key).copied().unwrap_or_default())
}

/// Set usage directly (e.g., after a rescan).
pub fn set_usage(dataset: &str, key: QuotaKey, usage: QuotaUsage) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    let ds = store.get_or_create(dataset);

    ds.usage.insert(key, usage);

    // Update quota if it exists
    if let Some(quota) = ds.quotas.get_mut(&key) {
        quota.usage = usage;
    }

    Ok(())
}

/// Update usage (add or subtract delta).
pub fn update_usage(
    dataset: &str,
    uid: u32,
    gid: u32,
    bytes_delta: i64,
    inodes_delta: i64,
    now: u64,
) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    let ds = store.get_or_create(dataset);

    // Update user usage
    let user_key = QuotaKey::user(uid);
    update_key_usage(ds, user_key, bytes_delta, inodes_delta, now);

    // Update group usage
    let group_key = QuotaKey::group(gid);
    update_key_usage(ds, group_key, bytes_delta, inodes_delta, now);

    Ok(())
}

/// Helper to update usage for a specific key.
fn update_key_usage(
    ds: &mut DatasetQuotas,
    key: QuotaKey,
    bytes_delta: i64,
    inodes_delta: i64,
    now: u64,
) {
    let usage = ds.usage.entry(key).or_default();

    if bytes_delta >= 0 {
        usage.add_bytes(bytes_delta as u64);
    } else {
        usage.sub_bytes((-bytes_delta) as u64);
    }

    if inodes_delta >= 0 {
        for _ in 0..inodes_delta {
            usage.add_inode();
        }
    } else {
        for _ in 0..(-inodes_delta) {
            usage.sub_inode();
        }
    }

    // Update quota if exists
    if let Some(quota) = ds.quotas.get_mut(&key) {
        quota.usage = *usage;
        quota.update_exceeded_timestamps(now);
    }
}

/// Check if a write operation is allowed.
pub fn check_write(
    dataset: &str,
    uid: u32,
    gid: u32,
    size: u64,
    now: u64,
) -> QuotaResult<QuotaCheckResult> {
    let store = QUOTA_STORE.lock();
    let ds = match store.get(dataset) {
        Some(ds) => ds,
        None => return Ok(QuotaCheckResult::ok()), // No quotas configured
    };

    let mut result = QuotaCheckResult::ok();
    let mut min_bytes = u64::MAX;

    // Check user quota
    let user_key = QuotaKey::user(uid);
    if let Some(quota) = ds.quotas.get(&user_key) {
        let user_usage = ds.usage.get(&user_key).copied().unwrap_or_default();
        let status = check_quota_limits(quota, &user_usage, size, 0, now);
        result.user_status = Some(status);

        if let Some(q) = ds.quotas.get(&user_key) {
            let remaining = q.limits.hard_bytes.saturating_sub(user_usage.bytes_used);
            min_bytes = min_bytes.min(remaining);
        }

        if status.is_error() {
            result.status = status;
        } else if status.is_warning() && result.status == QuotaStatus::Ok {
            result.status = status;
            result.add_warning(alloc::format!("user {} approaching quota limit", uid));
        }
    }

    // Check group quota
    let group_key = QuotaKey::group(gid);
    if let Some(quota) = ds.quotas.get(&group_key) {
        let group_usage = ds.usage.get(&group_key).copied().unwrap_or_default();
        let status = check_quota_limits(quota, &group_usage, size, 0, now);
        result.group_status = Some(status);

        if let Some(q) = ds.quotas.get(&group_key) {
            let remaining = q.limits.hard_bytes.saturating_sub(group_usage.bytes_used);
            min_bytes = min_bytes.min(remaining);
        }

        if status.is_error() {
            result.status = status;
        } else if status.is_warning() && result.status == QuotaStatus::Ok {
            result.status = status;
            result.add_warning(alloc::format!("group {} approaching quota limit", gid));
        }
    }

    result.bytes_allowed = min_bytes;
    Ok(result)
}

/// Check if creating a file/directory is allowed.
pub fn check_create(dataset: &str, uid: u32, gid: u32, now: u64) -> QuotaResult<QuotaCheckResult> {
    let store = QUOTA_STORE.lock();
    let ds = match store.get(dataset) {
        Some(ds) => ds,
        None => return Ok(QuotaCheckResult::ok()),
    };

    let mut result = QuotaCheckResult::ok();
    let mut min_inodes = u64::MAX;

    // Helper to update result status based on quota check
    let update_status = |result: &mut QuotaCheckResult, status: QuotaStatus| {
        let should_update =
            status.is_error() || (status.is_warning() && result.status == QuotaStatus::Ok);
        if should_update {
            result.status = status;
        }
    };

    // Check user quota
    let user_key = QuotaKey::user(uid);
    if let Some(quota) = ds.quotas.get(&user_key) {
        let user_usage = ds.usage.get(&user_key).copied().unwrap_or_default();
        let status = check_quota_limits(quota, &user_usage, 0, 1, now);
        result.user_status = Some(status);

        if let Some(q) = ds.quotas.get(&user_key) {
            let remaining = q.limits.hard_inodes.saturating_sub(user_usage.inodes_used);
            min_inodes = min_inodes.min(remaining);
        }

        update_status(&mut result, status);
    }

    // Check group quota
    let group_key = QuotaKey::group(gid);
    if let Some(quota) = ds.quotas.get(&group_key) {
        let group_usage = ds.usage.get(&group_key).copied().unwrap_or_default();
        let status = check_quota_limits(quota, &group_usage, 0, 1, now);
        result.group_status = Some(status);

        if let Some(q) = ds.quotas.get(&group_key) {
            let remaining = q.limits.hard_inodes.saturating_sub(group_usage.inodes_used);
            min_inodes = min_inodes.min(remaining);
        }

        update_status(&mut result, status);
    }

    result.inodes_allowed = min_inodes;
    Ok(result)
}

/// Check quota limits and return status.
fn check_quota_limits(
    quota: &Quota,
    usage: &QuotaUsage,
    bytes_add: u64,
    inodes_add: u64,
    now: u64,
) -> QuotaStatus {
    let new_bytes = usage.bytes_used.saturating_add(bytes_add);
    let new_inodes = usage.inodes_used.saturating_add(inodes_add);

    // Check hard limits
    if quota.limits.hard_bytes > 0 && new_bytes > quota.limits.hard_bytes {
        return QuotaStatus::HardLimitExceeded;
    }
    if quota.limits.hard_inodes > 0 && new_inodes > quota.limits.hard_inodes {
        return QuotaStatus::HardLimitExceeded;
    }

    // Check soft limits with grace period
    let bytes_over_soft = quota.limits.soft_bytes > 0 && new_bytes > quota.limits.soft_bytes;
    let inodes_over_soft = quota.limits.soft_inodes > 0 && new_inodes > quota.limits.soft_inodes;

    if bytes_over_soft {
        if usage.is_bytes_grace_expired(now, quota.limits.grace_period) {
            return QuotaStatus::SoftLimitExceeded;
        }
        return QuotaStatus::SoftLimitWarning;
    }

    if inodes_over_soft {
        if usage.is_inodes_grace_expired(now, quota.limits.grace_period) {
            return QuotaStatus::SoftLimitExceeded;
        }
        return QuotaStatus::SoftLimitWarning;
    }

    QuotaStatus::Ok
}

/// Generate a quota report for a dataset.
pub fn generate_report(dataset: &str, now: u64) -> QuotaResult<QuotaReport> {
    let store = QUOTA_STORE.lock();
    let ds = store
        .get(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    let mut report = QuotaReport {
        dataset: dataset.to_string(),
        total_entries: ds.quotas.len() as u32,
        over_soft: 0,
        over_hard: 0,
        grace_expired: 0,
        entries: Vec::new(),
    };

    for quota in ds.quotas.values() {
        let entry = QuotaReportEntry::from_quota(quota, now);

        match entry.status {
            QuotaStatus::HardLimitExceeded => report.over_hard += 1,
            QuotaStatus::SoftLimitExceeded => {
                report.over_soft += 1;
                report.grace_expired += 1;
            }
            QuotaStatus::SoftLimitWarning => report.over_soft += 1,
            QuotaStatus::Ok => {}
        }

        report.entries.push(entry);
    }

    Ok(report)
}

/// List all quotas for a dataset.
pub fn list_quotas(dataset: &str) -> QuotaResult<Vec<Quota>> {
    let store = QUOTA_STORE.lock();
    let ds = store
        .get(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    Ok(ds.quotas.values().cloned().collect())
}

/// List all quotas of a specific type.
pub fn list_quotas_by_type(dataset: &str, quota_type: QuotaType) -> QuotaResult<Vec<Quota>> {
    let store = QUOTA_STORE.lock();
    let ds = store
        .get(dataset)
        .ok_or_else(|| QuotaError::DatasetNotFound(dataset.into()))?;

    Ok(ds
        .quotas
        .values()
        .filter(|q| q.key.quota_type == quota_type)
        .cloned()
        .collect())
}

/// Clear all quotas for a dataset.
pub fn clear_quotas(dataset: &str) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    store.datasets.remove(dataset);
    Ok(())
}

/// Reset usage tracking for a dataset (e.g., before rescan).
pub fn reset_usage(dataset: &str) -> QuotaResult<()> {
    let mut store = QUOTA_STORE.lock();
    if let Some(ds) = store.get_mut(dataset) {
        ds.usage.clear();
        for quota in ds.quotas.values_mut() {
            quota.usage = QuotaUsage::zero();
        }
    }
    Ok(())
}

/// Get dataset names with quotas configured.
pub fn list_datasets() -> Vec<String> {
    let store = QUOTA_STORE.lock();
    store.datasets.keys().cloned().collect()
}

/// Check if a dataset has any quotas configured.
pub fn has_quotas(dataset: &str) -> bool {
    let store = QUOTA_STORE.lock();
    store
        .get(dataset)
        .map(|ds| !ds.quotas.is_empty())
        .unwrap_or(false)
}

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

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

    fn clean_test_dataset(name: &str) {
        let _ = clear_quotas(name);
    }

    #[test]
    fn test_set_get_quota() {
        let dataset = "test_set_get";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1000);
        let limits = QuotaLimits::bytes(1_000_000, 2_000_000);

        set_quota(dataset, key, limits).unwrap();
        let quota = get_quota(dataset, key).unwrap();

        assert_eq!(quota.limits.soft_bytes, 1_000_000);
        assert_eq!(quota.limits.hard_bytes, 2_000_000);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_remove_quota() {
        let dataset = "test_remove";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1001);
        set_quota(dataset, key, QuotaLimits::bytes(100, 200)).unwrap();
        remove_quota(dataset, key).unwrap();

        assert!(get_quota(dataset, key).is_err());

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_update_usage() {
        let dataset = "test_usage";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1002);
        set_quota(dataset, key, QuotaLimits::bytes(1000, 2000)).unwrap();

        update_usage(dataset, 1002, 100, 500, 1, 0).unwrap();

        let usage = get_usage(dataset, key).unwrap();
        assert_eq!(usage.bytes_used, 500);
        assert_eq!(usage.inodes_used, 1);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_check_write_ok() {
        let dataset = "test_check_ok";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1003);
        set_quota(dataset, key, QuotaLimits::bytes(1000, 2000)).unwrap();

        let result = check_write(dataset, 1003, 100, 500, 0).unwrap();
        assert!(result.is_allowed());
        assert_eq!(result.status, QuotaStatus::Ok);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_check_write_hard_limit() {
        let dataset = "test_check_hard";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1004);
        set_quota(dataset, key, QuotaLimits::bytes(1000, 2000)).unwrap();
        set_usage(dataset, key, QuotaUsage::new(1500, 0)).unwrap();

        // Try to write more than remaining
        let result = check_write(dataset, 1004, 100, 1000, 0).unwrap();
        assert!(!result.is_allowed());
        assert_eq!(result.status, QuotaStatus::HardLimitExceeded);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_check_write_soft_limit() {
        let dataset = "test_check_soft";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1005);
        set_quota(dataset, key, QuotaLimits::bytes(1000, 2000)).unwrap();
        set_usage(dataset, key, QuotaUsage::new(800, 0)).unwrap();

        // Writing 300 goes over soft (1000) but under hard (2000)
        let result = check_write(dataset, 1005, 100, 300, 0).unwrap();
        assert!(result.is_allowed());
        assert_eq!(result.status, QuotaStatus::SoftLimitWarning);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_generate_report() {
        let dataset = "test_report";
        clean_test_dataset(dataset);

        set_quota(dataset, QuotaKey::user(1), QuotaLimits::bytes(100, 200)).unwrap();
        set_quota(dataset, QuotaKey::user(2), QuotaLimits::bytes(100, 200)).unwrap();
        set_quota(dataset, QuotaKey::group(1), QuotaLimits::bytes(500, 1000)).unwrap();

        let report = generate_report(dataset, 0).unwrap();
        assert_eq!(report.total_entries, 3);
        assert_eq!(report.entries.len(), 3);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_list_quotas() {
        let dataset = "test_list";
        clean_test_dataset(dataset);

        set_quota(dataset, QuotaKey::user(1), QuotaLimits::bytes(100, 200)).unwrap();
        set_quota(dataset, QuotaKey::user(2), QuotaLimits::bytes(100, 200)).unwrap();
        set_quota(dataset, QuotaKey::group(1), QuotaLimits::bytes(500, 1000)).unwrap();

        let all = list_quotas(dataset).unwrap();
        assert_eq!(all.len(), 3);

        let users = list_quotas_by_type(dataset, QuotaType::User).unwrap();
        assert_eq!(users.len(), 2);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_reset_usage() {
        let dataset = "test_reset";
        clean_test_dataset(dataset);

        let key = QuotaKey::user(1006);
        set_quota(dataset, key, QuotaLimits::bytes(1000, 2000)).unwrap();
        set_usage(dataset, key, QuotaUsage::new(500, 5)).unwrap();

        reset_usage(dataset).unwrap();

        let usage = get_usage(dataset, key).unwrap();
        assert_eq!(usage.bytes_used, 0);

        clean_test_dataset(dataset);
    }

    #[test]
    fn test_has_quotas() {
        let dataset = "test_has";
        clean_test_dataset(dataset);

        assert!(!has_quotas(dataset));

        set_quota(dataset, QuotaKey::user(1), QuotaLimits::bytes(100, 200)).unwrap();
        assert!(has_quotas(dataset));

        clean_test_dataset(dataset);
    }
}