use std::{
    fs::File,
    io,
    path::{Path, PathBuf},
};

use serde::Serialize;
use thiserror::Error;

use crate::measure::{
    parse::{SelectiveStatFile, StatFileBuilder, U64File},
    v2::{
        line_index::LineIndex,
        settings::{EnabledKeys, EnabledKeysError},
    },
};

/// Collects measurements from `memory.current`.
pub struct MemoryCurrentCollector {
    file: U64File,
}

/// Collects measurements from `memory.stat`.
pub struct MemoryStatCollector {
    stat_file: SelectiveStatFile,
    mapping: MemoryStatMapping,
}

/// Represents the measurements extracted from the `memory.stat` file.
#[derive(Default)]
pub struct MemoryStats {
    pub anon: Option<u64>,
    pub file: Option<u64>,
    pub kernel_stack: Option<u64>,
    pub page_tables: Option<u64>,
    // could be extended to manage other memory.stat measurements
}

#[derive(Default)]
struct MemoryStatMapping {
    anon: LineIndex,
    file: LineIndex,
    kernel_stack: LineIndex,
    page_tables: LineIndex,
}

#[derive(Debug, Serialize)]
pub struct MemoryStatCollectorSettings {
    pub anon: bool,
    pub file: bool,
    pub kernel_stack: bool,
    #[serde(rename = "pagetables")]
    pub page_tables: bool,
}

impl EnabledKeys for MemoryStatCollectorSettings {}

impl Default for MemoryStatCollectorSettings {
    fn default() -> Self {
        Self {
            anon: true,
            file: true,
            kernel_stack: true,
            page_tables: true,
        }
    }
}

impl MemoryCurrentCollector {
    pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let file = U64File::open(path)?;
        Ok(Self { file })
    }

    /// 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<u64> {
        // SAFETY: the content is generated by the kernel and is always valid ASCII (hence valid UTF-8)
        unsafe { self.file.read(io_buf) }
    }
}

#[derive(Debug, Error)]
pub enum CollectorCreationError {
    #[error("failed to create collector: error on {1}")]
    Io(#[source] io::Error, PathBuf),
    #[error("failed to create collector: bad settings")]
    Settings(#[from] EnabledKeysError),
}

impl MemoryStatCollector {
    pub fn new<P: AsRef<Path>>(
        path: P,
        settings: MemoryStatCollectorSettings,
        io_buf: &mut Vec<u8>,
    ) -> Result<Self, CollectorCreationError> {
        let path = path.as_ref();
        let file = File::open(path).map_err(|e| CollectorCreationError::Io(e, path.into()))?;

        let keys = settings.enabled_keys()?;
        let (stat_file, stat_mapping) = StatFileBuilder::new(file, &keys)
            .build(io_buf.as_mut())
            .map_err(|e| CollectorCreationError::Io(e, path.into()))?;

        let mut mapping = MemoryStatMapping::default();
        if let Some(i) = stat_mapping.line_index("anon") {
            mapping.anon = i.into();
        }
        if let Some(i) = stat_mapping.line_index("file") {
            mapping.file = i.into();
        }
        if let Some(i) = stat_mapping.line_index("kernel_stack") {
            mapping.kernel_stack = i.into();
        }
        if let Some(i) = stat_mapping.line_index("pagetables") {
            mapping.page_tables = i.into();
        }

        if !stat_mapping.keys_not_found().is_empty() {
            log::warn!(
                "keys not found in {}: {}",
                path.display(),
                stat_mapping.keys_not_found().join(", ")
            )
        }

        Ok(Self { stat_file, mapping })
    }

    /// 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<MemoryStats> {
        let mut res = MemoryStats::default();
        unsafe {
            self.stat_file.read(io_buf, |i, _k, v| match i {
                i if i == self.mapping.anon.0 => {
                    res.anon = Some(v);
                }
                i if i == self.mapping.file.0 => {
                    res.file = Some(v);
                }
                i if i == self.mapping.kernel_stack.0 => {
                    res.kernel_stack = Some(v);
                }
                i if i == self.mapping.page_tables.0 => {
                    res.page_tables = Some(v);
                }
                _ => (),
            })
        }?;
        Ok(res)
    }
}

#[cfg(test)]
mod tests {
    use std::io::ErrorKind;

    use crate::measure::v2::{
        memory::{MemoryCurrentCollector, MemoryStatCollector, MemoryStatCollectorSettings},
        mock::{MemoryStatMock, MockFileCgroupKV},
    };

    #[test]
    fn collect_memory_stat() -> anyhow::Result<()> {
        let mut tmp = tempfile::NamedTempFile::new()?;

        let mock = MemoryStatMock {
            anon: 63,
            file: 12,
            kernel_stack: 123,
            pagetables: 42,
            ..Default::default()
        };
        mock.write_to_file(tmp.as_file_mut())?;

        let mut io_buf = Vec::new();
        let mut collector = MemoryStatCollector::new(
            tmp.path(),
            MemoryStatCollectorSettings {
                anon: true,
                file: false,
                kernel_stack: true,
                page_tables: true,
            },
            io_buf.as_mut(),
        )?;

        let memory_stats = collector.measure(io_buf.as_mut())?;
        assert_eq!(memory_stats.anon, Some(63));
        assert_eq!(memory_stats.file, None);
        assert_eq!(memory_stats.kernel_stack, Some(123));
        assert_eq!(memory_stats.page_tables, Some(42));
        Ok(())
    }

    #[test]
    fn collect_memory_stat_empty() -> anyhow::Result<()> {
        let tmp = tempfile::NamedTempFile::new()?;

        // completely empty file
        std::fs::write(tmp.path(), "")?;
        let mut io_buf = Vec::new();
        let mut collector = MemoryStatCollector::new(
            tmp.path(),
            MemoryStatCollectorSettings {
                anon: true,
                file: true,
                kernel_stack: true,
                page_tables: true,
            },
            io_buf.as_mut(),
        )?;

        let memory_stats = collector.measure(io_buf.as_mut())?;
        assert_eq!(memory_stats.anon, None);
        assert_eq!(memory_stats.file, None);
        assert_eq!(memory_stats.kernel_stack, None);
        assert_eq!(memory_stats.page_tables, None);

        // empty lines
        std::fs::write(tmp.path(), "\n\n\n")?;
        let memory_stats = collector.measure(io_buf.as_mut())?;
        assert_eq!(memory_stats.anon, None);
        assert_eq!(memory_stats.file, None);
        assert_eq!(memory_stats.kernel_stack, None);
        assert_eq!(memory_stats.page_tables, None);
        Ok(())
    }

    #[test]
    fn collect_memory_current() -> anyhow::Result<()> {
        let tmp = tempfile::NamedTempFile::new()?;

        let mut io_buf = Vec::new();
        let mut collector = MemoryCurrentCollector::new(tmp.path())?;

        std::fs::write(tmp.path(), "1999")?;
        let res = collector.measure(io_buf.as_mut())?;
        assert_eq!(res, 1999);

        std::fs::write(tmp.path(), "0")?;
        let res = collector.measure(io_buf.as_mut())?;
        assert_eq!(res, 0);

        std::fs::write(tmp.path(), "42\n")?;
        let res = collector.measure(io_buf.as_mut())?;
        assert_eq!(res, 42);
        Ok(())
    }

    #[test]
    fn bad_memory_current() -> anyhow::Result<()> {
        let tmp = tempfile::NamedTempFile::new()?;

        let mut io_buf = Vec::new();
        let mut collector = MemoryCurrentCollector::new(tmp.path())?;

        std::fs::write(tmp.path(), "")?;
        let res = collector.measure(io_buf.as_mut()).expect_err("should fail");
        assert_eq!(res.kind(), ErrorKind::InvalidData);

        std::fs::write(tmp.path(), "\n\n")?;
        let res = collector.measure(io_buf.as_mut()).expect_err("should fail");
        assert_eq!(res.kind(), ErrorKind::InvalidData);
        Ok(())
    }
}
