//
// Syd: rock-solid application kernel
// src/utils/syd-sys.rs: Calculate the memory usage of a given process or the parent process.
//
// Copyright (c) 2024, 2025 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

use std::process::ExitCode;

use nix::{errno::Errno, libc::pid_t, unistd::Pid};
use syd::{
    human_size,
    proc::{proc_mem, proc_smaps, proc_statm},
};

// Set global allocator to GrapheneOS allocator.
#[cfg(all(
    not(coverage),
    not(feature = "prof"),
    not(target_os = "android"),
    target_page_size_4k,
    target_pointer_width = "64"
))]
#[global_allocator]
static GLOBAL: hardened_malloc::HardenedMalloc = hardened_malloc::HardenedMalloc;

// Set global allocator to tcmalloc if profiling is enabled.
#[cfg(feature = "prof")]
#[global_allocator]
static GLOBAL: tcmalloc::TCMalloc = tcmalloc::TCMalloc;

syd::main! {
    use lexopt::prelude::*;

    syd::set_sigpipe_dfl()?;

    // Configure syd::proc.
    syd::config::proc_init_simple()?;

    // Parse CLI options.
    let mut opt_smaps = false; // -s (use the slow mode).
    let mut opt_human = false; // -H
    let mut opt_is_vm = false; // -V
    let mut opt_pid = None;

    let mut parser = lexopt::Parser::from_env();
    while let Some(arg) = parser.next()? {
        match arg {
            Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Short('H') => opt_human = true,
            Short('V') => opt_is_vm = true,
            Short('s') => opt_smaps = true,
            Value(pid) if opt_pid.is_none() => {
                opt_pid = Some(pid.parse::<pid_t>()?);
            }
            _ => return Err(arg.unexpected().into()),
        }
    }

    let pid = match opt_pid {
        None => Pid::parent(),
        Some(pid) => Pid::from_raw(pid),
    };

    let size = if opt_is_vm {
        match proc_statm(pid) {
            Ok(statm) => statm.size.saturating_mul(*syd::config::PAGE_SIZE),
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }
    } else if opt_smaps {
        match proc_mem_smaps(pid) {
            Ok(size) => size,
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }
    } else {
        match proc_mem(pid) {
            Ok(size) => size,
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }
    };

    if opt_human {
        println!("{}", human_size(size as usize));
    } else {
        println!("{size}");
    }

    Ok(ExitCode::SUCCESS)
}

fn help() {
    println!("Usage: syd-mem [-HV] [pid]");
    println!("Calculate the memory usage of a given process or the parent process and exit.");
    println!("-H    Print human-formatted size");
    println!("-V    Print virtual memory size");
}

/// Calculates process memory usage (slowly) using proc_pid_smaps(5).
///
/// This function iterates through proc_pid_smaps(5), and therefore it's less efficient than
/// `syd::proc::proc_mem` which calculates memory usage using the file _proc_pid_smaps_rollup(5)
/// in one go. See: https://www.kernel.org/doc/Documentation/ABI/testing/procfs-smaps_rollup
fn proc_mem_smaps(pid: Pid) -> Result<u64, Errno> {
    proc_smaps(pid).map(|maps| {
        let mut sum = 0u64;
        for map in &maps {
            for key in ["Pss", "Private_Dirty", "Shared_Dirty"] {
                let val = map.0.extension.map.get(key).copied().unwrap_or(0);
                sum = sum.saturating_add(val);
            }
        }
        sum
    })
}
