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

use std::process::ExitCode;

use getargs::{Opt, Options};
use nix::{
    libc::pid_t,
    unistd::{sysconf, Pid, SysconfVar},
};
use once_cell::sync::Lazy;
use procfs::{
    process::{MMapPath, Process},
    ProcError,
};
use syd::{err::SydResult, human_size};

// Set global allocator to mimalloc.
#[global_allocator]
static GLOBAL: mimalloc2_rust::GlobalMiMalloc = mimalloc2_rust::GlobalMiMalloc;

// System page size
static PAGE_SIZE: Lazy<u64> = Lazy::new(|| {
    sysconf(SysconfVar::PAGE_SIZE)
        .unwrap_or(Some(4096))
        .unwrap_or(4096) as u64
});

fn main() -> SydResult<ExitCode> {
    syd::set_sigpipe_dfl()?;

    // Configure mimalloc.
    syd::configure_mimalloc();

    let args = std::env::args().skip(1).collect::<Vec<_>>();
    let args: &'static Vec<String> = Box::leak(Box::new(args));
    let mut opts = Options::new(args.iter().map(String::as_str));
    let mut human = false; // -H
    let mut is_vm = false; // -V

    while let Some(opt) = opts.next_opt()? {
        match opt {
            Opt::Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Opt::Short('H') => {
                human = true;
            }
            Opt::Short('V') => {
                is_vm = true;
            }
            Opt::Short(opt) => {
                eprintln!("Unknown option: -{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
            Opt::Long(opt) => {
                eprintln!("Unknown option: --{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
        }
    }

    let pid = match opts.positionals().next() {
        None => {
            // Find parent pid
            match Process::new(Pid::this().as_raw()).and_then(|proc| proc.stat()) {
                Ok(stat) => stat.ppid as pid_t,
                Err(error) => {
                    eprintln!("syd-mem: {error}");
                    return Ok(ExitCode::FAILURE);
                }
            }
        }
        Some(pid) => match pid.parse::<pid_t>() {
            Ok(pid) => pid,
            Err(error) => {
                eprintln!("syd-mem: {error}");
                help();
                return Ok(ExitCode::FAILURE);
            }
        },
    };

    let size = if is_vm {
        match Process::new(pid).and_then(|proc| proc.statm()) {
            Ok(statm) => statm.size.saturating_mul(*PAGE_SIZE),
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }
    } else {
        let proc = match Process::new(pid) {
            Ok(proc) => proc,
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        };
        match proc_mem(&proc) {
            Ok(size) => size,
            Err(error) => {
                eprintln!("syd-mem: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }
    };

    if 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.
///
/// This function uses the `procfs` crate to obtain detailed memory maps
/// from `/proc/[pid]/smaps`. It sums multiple memory usage values reported in these maps
/// to calculate a more comprehensive total memory usage.
///
/// # Arguments
///
/// * `process` - `Process` instance representing the process.
///
/// # Returns
///
/// This function returns a `Result<u64, Errno>`.
///
/// # Errors
///
/// This function returns an error if it fails to retrieve the process's memory maps,
/// typically due to insufficient permissions or an invalid process ID.
fn proc_mem(process: &Process) -> Result<u64, ProcError> {
    process.smaps().map(|maps| {
        let mut total_size: u64 = 0;
        for map in &maps.0 {
            match &map.pathname {
                MMapPath::Path(_) | MMapPath::Anonymous | MMapPath::Stack | MMapPath::Other(_) => {
                    let pss = map.extension.map.get("Pss").copied().unwrap_or(0);
                    let private_dirty =
                        map.extension.map.get("Private_Dirty").copied().unwrap_or(0);
                    let shared_dirty = map.extension.map.get("Shared_Dirty").copied().unwrap_or(0);

                    total_size = total_size.saturating_add(
                        pss.saturating_add(private_dirty)
                            .saturating_add(shared_dirty),
                    );
                }
                _ => (),
            }
        }

        total_size
    })
}
