// Copyright 2026 PingCAP, Inc.

use std::{
    path::PathBuf,
    sync::atomic::{AtomicU64, Ordering},
};

use tikv_util::{error, sys::disk, time::Duration};

/// Collecting storage stats.
/// Uses the same interval of reporing store heartbeat for it to reduce
/// frequently accessing statistics on engines.
static DEFAULT_STORAGE_STATS_INTERVAL: AtomicU64 =
    AtomicU64::new(Duration::from_secs(10).as_millis() as u64);

// Only for tests.
pub fn reset_update_storage_stats_interval(duration: Duration) {
    DEFAULT_STORAGE_STATS_INTERVAL.store(duration.as_millis() as u64, Ordering::Relaxed);
}

pub fn get_update_storage_stats_interval() -> Duration {
    Duration::from_millis(DEFAULT_STORAGE_STATS_INTERVAL.load(Ordering::Relaxed))
}

fn calculate_disk_usage(a: disk::DiskUsage, b: disk::DiskUsage) -> disk::DiskUsage {
    match (a, b) {
        (disk::DiskUsage::AlreadyFull, _) => disk::DiskUsage::AlreadyFull,
        (_, disk::DiskUsage::AlreadyFull) => disk::DiskUsage::AlreadyFull,
        (disk::DiskUsage::AlmostFull, _) => disk::DiskUsage::AlmostFull,
        (_, disk::DiskUsage::AlmostFull) => disk::DiskUsage::AlmostFull,
        (disk::DiskUsage::Normal, disk::DiskUsage::Normal) => disk::DiskUsage::Normal,
    }
}

/// A checker to inspect the disk usage of kv engine and raft engine.
/// The caller should call `inspect` periodically to get the disk usage status
/// manually.
#[derive(Clone)]
pub struct DiskUsageChecker {
    /// The data paths of kv engine.
    kvdb_paths: Vec<PathBuf>,
    /// The path of raft engine.
    raft_path: PathBuf,
    /// The path of auxiliary directory of raft engine if specified.
    raft_auxiliary_path: Option<PathBuf>,
    /// Whether the main directory of raft engine is separated from kv engine.
    separated_raft_mount_path: bool,
    /// Whether the auxiliary directory of raft engine is separated from kv
    /// engine.
    separated_raft_auxiliary_mount_path: bool,
    /// Whether the auxiliary directory of raft engine is both separated from
    /// the main directory of raft engine and kv engine.
    separated_raft_auxiliary_and_kvdb_mount_path: bool,
    /// The threshold of disk usage of kv engine to trigger the almost full
    /// status.
    kvdb_almost_full_thd: u64,
    /// The threshold of disk usage of raft engine to trigger the almost full
    /// status.
    raft_almost_full_thd: u64,
    /// The specified disk capacity for the whole disk.
    config_disk_capacity: u64,
}

impl DiskUsageChecker {
    pub fn new(
        kvdb_paths: Vec<PathBuf>,
        raft_path: PathBuf,
        raft_auxiliary_path: Option<PathBuf>,
        separated_raft_mount_path: bool,
        separated_raft_auxiliary_mount_path: bool,
        separated_raft_auxiliary_and_kvdb_mount_path: bool,
        kvdb_almost_full_thd: u64,
        raft_almost_full_thd: u64,
        config_disk_capacity: u64,
    ) -> Self {
        DiskUsageChecker {
            kvdb_paths,
            raft_path,
            raft_auxiliary_path,
            separated_raft_mount_path,
            separated_raft_auxiliary_mount_path,
            separated_raft_auxiliary_and_kvdb_mount_path,
            kvdb_almost_full_thd,
            raft_almost_full_thd,
            config_disk_capacity,
        }
    }

    /// Inspect the disk usage of kv engine and raft engine.
    ///
    /// Returns the disk usage status of the whole disk, kv engine and raft
    /// engine, the whole disk capacity and available size.
    ///
    /// Returns `None` when any required disk stat cannot be collected.
    pub fn inspect(
        &self,
    ) -> Option<(
        disk::DiskUsage, // whole disk status
        disk::DiskUsage, // kvdb disk status
        disk::DiskUsage, // raft disk status
        u64,             // whole capacity
        u64,             // whole available
    )> {
        // By default, the almost full threshold of kv engine is half of the
        // configured value.
        let kvdb_already_full_thd = self.kvdb_almost_full_thd / 2;
        let raft_already_full_thd = self.raft_almost_full_thd / 2;
        // Check the disk space of raft engine.
        let raft_disk_status = {
            if !self.separated_raft_mount_path || self.raft_almost_full_thd == 0 {
                disk::DiskUsage::Normal
            } else {
                let (_, raft_disk_avail) = match disk::get_disks_stats(std::slice::from_ref(
                    &self.raft_path,
                )) {
                    Err(e) => {
                        error!(
                            "get disk stat for raft engine failed";
                            "raft_engine_path" => %self.raft_path.display(),
                            "err" => ?e
                        );
                        return None;
                    }
                    Ok((cap, avail)) => {
                        if !self.separated_raft_auxiliary_mount_path {
                            // If the auxiliary directory of raft engine is not separated from
                            // kv engine, returns u64::MAX to indicate that the disk space of
                            // the raft engine should not be checked.
                            (u64::MAX, u64::MAX)
                        } else if self.separated_raft_auxiliary_and_kvdb_mount_path {
                            // If the auxiliary directory of raft engine is separated from kv
                            // engine and the main directory of
                            // raft engine, the disk space of
                            // the auxiliary directory should be
                            // checked.
                            assert!(self.raft_auxiliary_path.is_some());
                            let (auxiliary_disk_cap, auxiliary_disk_avail) =
                                match disk::get_disks_stats(&[self
                                    .raft_auxiliary_path
                                    .as_ref()
                                    .unwrap()
                                    .to_path_buf()])
                                {
                                    Err(e) => {
                                        error!(
                                            "get auxiliary disk stat for raft engine failed";
                                            "raft_engine_path" => %self.raft_auxiliary_path.as_ref().unwrap().display(),
                                            "err" => ?e
                                        );
                                        return None;
                                    }
                                    Ok((total, avail)) => (total, avail),
                                };
                            (cap + auxiliary_disk_cap, avail + auxiliary_disk_avail)
                        } else {
                            (cap, avail)
                        }
                    }
                };
                if raft_disk_avail <= raft_already_full_thd {
                    disk::DiskUsage::AlreadyFull
                } else if raft_disk_avail <= self.raft_almost_full_thd {
                    disk::DiskUsage::AlmostFull
                } else {
                    disk::DiskUsage::Normal
                }
            }
        };
        // Check the disk space of kv engine.
        let (disk_cap, disk_avail) = match disk::get_disks_stats(&self.kvdb_paths) {
            Err(e) => {
                error!(
                    "get disk stat for kv store failed";
                    "kv_paths" => ?self.kvdb_paths,
                    "err" => ?e
                );
                return None;
            }
            Ok((total, avail)) => (total, avail),
        };
        let capacity = if self.config_disk_capacity == 0 || disk_cap < self.config_disk_capacity {
            disk_cap
        } else {
            self.config_disk_capacity
        };
        // Since the disk usage of kv engine consists of different types
        // of data, such as columnar files, sst files, IA segments, etc. So,
        // the available size of the kv engine should be the available size of
        // the local disk minus the used size of the kv engine.
        let available = disk_avail.saturating_sub(disk_cap.saturating_sub(capacity));
        let cur_kv_disk_status = if available <= kvdb_already_full_thd {
            disk::DiskUsage::AlreadyFull
        } else if available <= self.kvdb_almost_full_thd {
            disk::DiskUsage::AlmostFull
        } else {
            disk::DiskUsage::Normal
        };
        let cur_disk_status = calculate_disk_usage(raft_disk_status, cur_kv_disk_status);
        Some((
            cur_disk_status,
            cur_kv_disk_status,
            raft_disk_status,
            capacity,
            available,
        ))
    }
}