//! Measure cgroup v1 things.

use anyhow::Context;
use std::io::{self};

use crate::{Cgroup, CgroupHierarchy, measure::parse::U64File};

/// Collects cgroup v1 measurements.
pub struct V1Collector {
    cpuacct_usage: Option<U64File>,
    memory_usage: Option<U64File>,
}

/// The result of a cgroupv1 measurement operation.
pub struct V1Stats {
    pub cpuacct_usage: Option<u64>,
    pub memory_usage: Option<u64>,
}

impl V1Collector {
    /// Creates a new collector that gathers data from multiple cgroup hierarchies.
    ///
    /// # Available metrics
    /// The probe can only measure what is provided by the cgroup controllers attached to the given hierarchies.
    /// For instance, if one hierarchy has the `cpuacct` controller, and another one has the `memory` controller,
    /// the probe reads both `cpuacct.usage` and `memory.stat` (in their respective hierarchies).
    pub fn across_hierarchies(cgroup_path: &str, hierarchies: &[&CgroupHierarchy]) -> anyhow::Result<Self> {
        let mut cpuacct_usage = None;
        let mut memory_usage = None;
        for h in hierarchies {
            let cgroup_in_sysfs = h.cgroup_fs_path(cgroup_path);
            if h.available_controllers().iter().any(|c| c == "cpuacct") {
                let data_path = cgroup_in_sysfs.join("cpuacct.usage");
                let file = U64File::open(&data_path).with_context(|| format!("failed to open {data_path:?}"))?;
                cpuacct_usage = Some(file)
            }
            if h.available_controllers().iter().any(|c| c == "memory") {
                let data_path = cgroup_in_sysfs.join("memory.usage_in_bytes");
                let file = U64File::open(&data_path).with_context(|| format!("failed to open {data_path:?}"))?;
                memory_usage = Some(file)
            }
        }
        Ok(Self {
            cpuacct_usage,
            memory_usage,
        })
    }

    /// Creates a new `V1Collector` in a single cgroup v1 hierarchy.
    ///
    /// # Available metrics
    /// The collector can only measure what is provided by the cgroup controllers attached to the given hierarchy.
    pub fn in_single_hierarchy(cgroup: Cgroup) -> anyhow::Result<Self> {
        Self::across_hierarchies(cgroup.canonical_path(), &[cgroup.hierarchy()])
            .with_context(|| format!("collector creation failed for cgroup {}", cgroup.unique_name()))
    }

    /// Collects measurements from the underlying "file", using `io_buf` as an intermediary I/O buffer.
    pub fn measure(&mut self, io_buf: &mut Vec<u8>) -> io::Result<V1Stats> {
        // SAFETY: the content is generated by the kernel and is always valid ASCII (hence valid UTF-8)
        let cpuacct_usage = self
            .cpuacct_usage
            .as_mut()
            .map(|f| unsafe { f.read(io_buf) })
            .transpose()?;

        let memory_usage = self
            .memory_usage
            .as_mut()
            .map(|f| unsafe { f.read(io_buf) })
            .transpose()?;

        Ok(V1Stats {
            cpuacct_usage,
            memory_usage,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{fs::File, io::Write};

    use tempfile::tempdir;

    use crate::{CgroupHierarchy, CgroupVersion, measure::v1::V1Collector};

    #[test]
    fn test_across_hierarchies() -> anyhow::Result<()> {
        let root = tempdir().expect("Failed to create a temporary directory");
        // file 1
        let file_path = root.path().join("cpuacct.usage");
        let value: u64 = 15;
        let mut file1 = File::create(file_path)?;
        writeln!(file1, "{}", value)?;
        // file 2
        let file_path = root.path().join("memory.usage_in_bytes");
        let value: u64 = 19;
        let mut file2 = File::create(file_path)?;
        writeln!(file2, "{}", value)?;

        // let hierarchy = CgroupHierarchy::manually_unchecked_v1_named(root.path(), CgroupVersion::V1, vec!["cpu", "memory"]);
        // let hierarchy1 = CgroupHierarchy::manually_unchecked_v1_named(root.path(), "cpuacct".to_string());
        let hierarchy1 = CgroupHierarchy::manually_unchecked(root.path(), CgroupVersion::V2, vec!["cpuacct"]);
        let hierarchy2 = CgroupHierarchy::manually_unchecked(root.path(), CgroupVersion::V2, vec!["memory"]);

        let res = V1Collector::across_hierarchies("/", &[&hierarchy1, &hierarchy2]);
        let mut v1_collector = res.unwrap();
        let mut buf: Vec<u8> = vec![];

        assert!(v1_collector.cpuacct_usage.is_some());
        assert!(v1_collector.memory_usage.is_some());

        let res = v1_collector.measure(&mut buf);
        let collector = res.unwrap();
        assert_eq!(collector.cpuacct_usage, Some(15));
        assert_eq!(collector.memory_usage, Some(19));

        Ok(())
    }
}
