roundabout

use std::{

    collections::BTreeMap,

    path::Path,

    time::{Duration, SystemTime, UNIX_EPOCH},

};

use serde::{Deserialize, Serialize};

use crate::{sched::Need, store, units::Bytes};

const KEEP: usize = 1000;

#[derive(Clone, Copy, Debug, Serialize, Deserialize)]

pub struct Record {

    pub runs: u32,

    pub peak_rss: Bytes,

    /// ewma, so one freak run doesn't define the command forever

    pub wall: u64,

    pub cpu: u64,

    pub cores: u32,

    pub last: u64,

}

type History = BTreeMap<String, Record>;

pub fn key(root: &Path, cmd: &[String]) -> String {

    let mut k = format!("{}\t{}", root.display(), cmd.join(" "));

    k.truncate(k.floor_char_boundary(400));

    k

}

fn now() -> u64 {

    SystemTime::now()

        .duration_since(UNIX_EPOCH)

        .map_or(0, |d| d.as_secs())

}

pub fn record(key: String, peak_rss: Bytes, wall: Duration, cpu: Duration, cores: u32) {

    let _ = store::update("history", |h: &mut History| {

        let wall = wall.as_secs();

        h.entry(key)

            .and_modify(|r| {

                r.runs += 1;

                r.peak_rss = r.peak_rss.max(peak_rss);

                r.wall = (r.wall * 7 + wall * 3) / 10;

                r.cpu = cpu.as_secs();

                r.cores = cores;

                r.last = now();

            })

            .or_insert_with(|| Record {

                runs: 1,

                peak_rss,

                wall,

                cpu: cpu.as_secs(),

                cores,

                last: now(),

            });

        if h.len() > KEEP {

            let mut ages: Vec<_> = h.iter().map(|(k, r)| (r.last, k.clone())).collect();

            ages.sort();

            for (_, k) in ages.into_iter().take(h.len() - KEEP) {

                h.remove(&k);

            }

        }

    });

}

pub fn lookup(key: &str) -> Option<Record> {

    store::load::<History>("history").get(key).copied()

}

/// padded, since the agent that under-asks is the one that ooms the video.

/// nix builds happen in nix-daemon, so their rss never shows up in ours and

/// the base estimate stands

pub fn suggest(rec: &Record, base: Need, cmd: &[String]) -> Need {

    let via_daemon = cmd

        .iter()

        .any(|t| t.contains("nix") || t.contains("home-manager"));

    Need {

        ram: if via_daemon {

            base.ram

        } else {

            Bytes(rec.peak_rss.0 + rec.peak_rss.0 / 4).max(Bytes::mib(256))

        },

        cores: rec.cores.max(1),

        time: Duration::from_secs((rec.wall * 3 / 2).max(30)),

        ..base

    }

}