use clap::Parser;
use digit_group::FormatGroup;
use fastrand::i32 as ri32;
use memory_stats::memory_stats;
use std::iter::repeat_with;
use std::thread;
use std::time::{Duration, Instant};

#[derive(Parser, Debug)]
#[command(about = "Eats up that amount of memory and then quits on keypress", long_about = None)]
struct Args {
    #[clap(default_value_t = 16)]
    size_gb: u32,

    #[clap(
        short,
        long,
        default_value_t = 64,
        help = "Split into chunks of this size"
    )]
    block_size_mb: u32,

    #[clap(
        short,
        long,
        help = "If provided, wait this many seconds instead of expecting a keypress"
    )]
    wait: Option<u64>,

    #[clap(
        short,
        long,
        default_value_t = false,
        help = "Do not prompt if user is trying to allocate over 50% of their available RAM"
    )]
    ignore_warnings: bool,
}

fn main() {
    let args = Args::parse();
    // + 24 bytes will be used for each of these (assuming 64b arch).
    // Division by 4 is because we'll put 32bit integers in it to avoid OS memory compression.
    let block_size = args.block_size_mb * 1024 * 1024 / 4;

    // how many individual chunks of RAM to allocate on the heap
    // total blocks will be rounded down, so if you want to allocate 1GB = 1024MB RAM
    // with 100MB blocks, only 10x100MB will be used
    let total_blocks = args.size_gb * 1024 / args.block_size_mb;
    let total_size_mb = (total_blocks * args.block_size_mb) as u64;
    let total_size_b = total_size_mb * 1024 * 1024;
    let total_ram_bytes = system_memory::total();
    let allocation_ram_ratio = total_size_b as f64 / total_ram_bytes as f64;
    let ratio = (allocation_ram_ratio * 100.0) as u64;
    if !args.ignore_warnings && ratio >= 50 {
        if ratio >= 80 {
            println!("====================================================");
            println!("You're trying to allocate more than {ratio}% of your RAM!");
            println!("====================================================\n");
            println!(
                "This will most definitely cause system instability,\
            and if you're running on Windows, might require a reboot later\
            to properly clean up swap and cache everything normally again.\n"
            );
        } else {
            println!(
                "Warning! You're trying to allocate {ratio}% of your RAM.\
            Your system might slow down and misbehave.\n"
            );
        }
        println!(
            "You seem to have {}B RAM and you requested {} to be allocated.",
            total_ram_bytes.format_si('.'),
            total_size_b.format_si('.')
        );
        dont_disappear::enter_to_continue::custom_msg(
            "Press ENTER to continue this risky operation, or Ctrl-C to exit the program now.",
        );
    }

    println!(
        "Starting to allocate {total_size_mb}MB in {total_blocks} x {} MB blocks.",
        args.block_size_mb
    );
    let start = Instant::now();

    let mut all_blocks: Vec<Vec<i32>> = Vec::with_capacity(total_blocks as usize);

    for _ in 0..total_blocks {
        // let mut new_vec: Vec<f32> = Vec::with_capacity(block_size as usize);
        let new_vec: Vec<i32> = repeat_with(|| ri32(..)).take(block_size as usize).collect();
        all_blocks.push(new_vec);
    }

    let elapsed = start.elapsed();
    println!(
        "Allocation done in {} ms",
        (elapsed.as_millis() as u64).format_si('.')
    );
    let stats = memory_stats();
    if let Some(mem_stats) = stats {
        println!(
            "Actual \"physical\" memory occupied by the program: {} bytes",
            mem_stats.physical_mem.format_si('.')
        );
        println!(
            "Actual \"virtual\" memory occupied by the program: {} bytes",
            mem_stats.virtual_mem.format_si('.')
        );
        println!(
            "Physical / virtual size occupied by the program itself: {}/{} bytes",
            (mem_stats.physical_mem as u64 - total_size_b).format_si('.'),
            (mem_stats.virtual_mem as u64 - total_size_b).format_si('.')
        );
        // not really relevant unless we allow like 4kB blocks; even 32GB / 1MB
        // blocks still only uses 0.75MB, so the virtual size bloat comes from somewhere else?
        // let wide_pointer_size = size_of::<Vec<f32>>() as u64;
        // let pointer_usage = wide_pointer_size + (wide_pointer_size * total_blocks as u64);
        // println!("Bytes used by vectors themselves: {pointer_usage}B");
    }
    // println!("(1MB = 1024 x 1024 Bytes.)\n");
    if let Some(wait) = args.wait
        && wait > 0
    {
        println!("Sleeping {}s before quitting", wait);
        thread::sleep(Duration::from_secs(wait));
        println!("Freeing memory now.");
    } else {
        dont_disappear::any_key_to_continue::custom_msg("Press any key to exit");
    }
}
