roundabout

use std::{

    env,

    fmt::Write as _,

    os::unix::{net::UnixStream, process::CommandExt},

    process::{Command, ExitStatus},

    sync::{

        Arc,

        atomic::{AtomicBool, AtomicI32, Ordering},

    },

    thread,

    time::{Duration, Instant},

};

use crate::{

    artifacts, client,

    config::{Config, OverCap, RamOverrun},

    history, paths,

    proto::{self, Acquire, Event, Reader, Request},

    sched::Need,

    sys,

    units::{Amount, Bytes, Dur},

};

pub const EXIT_TIMEOUT: i32 = 124;

pub const EXIT_GAVE_UP: i32 = 75;

#[derive(Default)]

pub struct Args {

    pub size: Option<String>,

    pub ram: Option<Bytes>,

    pub disk: Option<Bytes>,

    pub cores: Option<Amount<u32>>,

    pub time: Option<Dur>,

    pub deadline: Option<Dur>,

    pub max_wait: Option<Dur>,

    pub session: Option<String>,

    pub exact: bool,

    pub cmd: Vec<String>,

}

static PGID: AtomicI32 = AtomicI32::new(0);

extern "C" fn forward(sig: libc::c_int) {

    let pgid = PGID.load(Ordering::Relaxed);

    if pgid > 1 {

        unsafe { libc::kill(-pgid, sig) };

    }

}

fn say(msg: &str) {

    eprintln!("rabt: {msg}");

}

pub fn need(cfg: &Config, a: &Args) -> Result<Need, String> {

    let base = cfg.preset(a.size.as_deref().unwrap_or(&cfg.enforce.default_size))?;

    Ok(Need {

        ram: a.ram.unwrap_or(base.ram),

        disk: a.disk.unwrap_or(base.disk),

        cores: a.cores.unwrap_or(base.cores).resolve(sys::cores()),

        time: a.time.unwrap_or(base.time).0,

    })

}

pub fn base(cfg: &Config) -> Result<Need, String> {

    let p = cfg.preset(&cfg.enforce.default_size)?;

    Ok(Need {

        ram: p.ram,

        disk: p.disk,

        cores: p.cores.resolve(sys::cores()),

        time: p.time.0,

    })

}

struct Grant {

    stream: UnixStream,

    lease: u64,

    cores: u32,

    ram: Bytes,

}

fn acquire(req: &Acquire, max_wait: Option<Dur>) -> Result<Grant, i32> {

    let started = Instant::now();

    // a daemon on its way out can accept us then hang up, so go round again

    for _ in 0..3 {

        let mut stream = client::connect_or_spawn().map_err(|e| {

            say(&format!("can't reach the daemon: {e}"));

            1

        })?;

        if proto::send(&mut stream, &Request::Acquire(req.clone())).is_err() {

            continue;

        }

        let _ = stream.set_read_timeout(Some(Duration::from_secs(1)));

        let mut reader = Reader::new(&stream).map_err(|_| 1)?;

        loop {

            match reader.recv::<Event>() {

                Ok(Some(Event::Granted {

                    lease,

                    cores,

                    ram,

                    heavy,

                })) => {

                    let bonus = cores.saturating_sub(req.need.cores);

                    say(&format!(

                        "lease {lease} go, cores {cores}{}{}",

                        if bonus > 0 {

                            format!(" ({bonus} spare)")

                        } else {

                            String::new()

                        },

                        if heavy { ", heavy" } else { "" }

                    ));

                    let _ = stream.set_read_timeout(None);

                    return Ok(Grant {

                        stream,

                        lease,

                        cores,

                        ram,

                    });

                }

                Ok(Some(Event::Queued {

                    lease,

                    position,

                    reason,

                })) => {

                    let why = reason.map(|r| format!(", {r}")).unwrap_or_default();

                    say(&format!("lease {lease} queued at #{position}{why}"));

                }

                Ok(Some(Event::Kill)) => {

                    say("killed while queued");

                    return Err(143);

                }

                Ok(Some(_)) => {}

                Err(e)

                    if matches!(

                        e.kind(),

                        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut

                    ) =>

                {

                    if max_wait.is_some_and(|m| started.elapsed() > m.0) {

                        say("gave up waiting");

                        return Err(EXIT_GAVE_UP);

                    }

                }

                Ok(None) | Err(_) => break,

            }

        }

    }

    say("daemon keeps hanging up");

    Err(1)

}

fn has_systemd_run() -> bool {

    env::var_os("PATH")

        .is_some_and(|p| env::split_paths(&p).any(|d| d.join("systemd-run").exists()))

}

fn command(cfg: &Config, a: &Args, need: &Need, grant: &Grant) -> Command {

    let mut argv = a.cmd.clone();

    if cfg.overrun.ram == RamOverrun::Throttle && cfg!(target_os = "linux") && has_systemd_run() {

        let high = Bytes((need.ram.0 as f64 * cfg.overrun.ram_slack) as u64);

        let mut wrapped: Vec<String> = [

            "systemd-run",

            "--user",

            "--scope",

            "--quiet",

            "--collect",

            "-p",

        ]

        .map(Into::into)

        .into();

        wrapped.push(format!("MemoryHigh={}", high.0));

        wrapped.push("--".into());

        wrapped.append(&mut argv);

        argv = wrapped;

    }

    let n = grant.cores.to_string();

    let mut cmd = Command::new(&argv[0]);

    cmd.args(&argv[1..])

        .process_group(0)

        .env("ROUNDABOUT_LEASE", grant.lease.to_string())

        .env("CARGO_BUILD_JOBS", &n)

        .env("CMAKE_BUILD_PARALLEL_LEVEL", &n);

    let make = env::var("MAKEFLAGS").unwrap_or_default();

    if !make.contains("-j") {

        cmd.env("MAKEFLAGS", format!("-j{n} {make}").trim_end());

    }

    if cfg.nix.inject {

        let per = cfg.nix.cores_per_job.max(1);

        let jobs = (grant.cores / per).max(1);

        let cores = grant.cores.min(per);

        let mut nix = env::var("NIX_CONFIG").unwrap_or_default();

        let _ = write!(nix, "\ncores = {cores}\nmax-jobs = {jobs}\n");

        cmd.env("NIX_CONFIG", nix);

    }

    #[cfg(target_os = "linux")]

    unsafe {

        cmd.pre_exec(|| {

            libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);

            Ok(())

        });

    }

    cmd

}

fn cpu_time() -> Duration {

    let mut ru = std::mem::MaybeUninit::<libc::rusage>::uninit();

    unsafe {

        if libc::getrusage(libc::RUSAGE_CHILDREN, ru.as_mut_ptr()) != 0 {

            return Duration::ZERO;

        }

        let ru = ru.assume_init();

        let tv = |t: libc::timeval| {

            Duration::from_secs(t.tv_sec as u64) + Duration::from_micros(t.tv_usec as u64)

        };

        tv(ru.ru_utime) + tv(ru.ru_stime)

    }

}

fn code(status: ExitStatus) -> i32 {

    use std::os::unix::process::ExitStatusExt;

    status

        .code()

        .or_else(|| status.signal().map(|s| 128 + s))

        .unwrap_or(1)

}

pub fn run(cfg: &Config, a: &Args) -> i32 {

    if a.cmd.is_empty() {

        say("nothing to run, put the command after --");

        return 2;

    }

    if a.size.is_none() && a.ram.is_none() {

        let sizes: Vec<_> = cfg.sizes.keys().map(String::as_str).collect();

        say(&format!(

            "how big is it? pass --size {} or --ram and friends. `rabt suggest -- <cmd>` has a guess",

            sizes.join("|")

        ));

        return 2;

    }

    let mut need = match need(cfg, a) {

        Ok(n) => n,

        Err(e) => {

            say(&e);

            return 2;

        }

    };

    let cwd = env::current_dir().unwrap_or_else(|_| "/".into());

    let root = paths::project_root(&cwd);

    let report = artifacts::check(cfg, Some(&root), false);

    if report.over() {

        let msg = report.describe(cfg);

        if cfg.artifacts.over_cap == OverCap::Block {

            say(&format!("{msg}. clean up first, `rabt artifacts gc --cap`"));

            return EXIT_GAVE_UP;

        }

        say(&format!(

            "{msg}. `rabt artifacts gc --cap` when you get a sec"

        ));

    }

    let req = Acquire {

        need,

        cwd,

        cmd: a.cmd.clone(),

        session: a.session.clone(),

        deadline: a.deadline,

        exact: a.exact,

    };

    let grant = match acquire(&req, a.max_wait.or(cfg.queue.max_wait)) {

        Ok(g) => g,

        Err(code) => return code,

    };

    need.ram = need.ram.max(grant.ram);

    let mut child = match command(cfg, a, &need, &grant).spawn() {

        Ok(c) => c,

        Err(e) => {

            say(&format!("couldn't start {}: {e}", a.cmd[0]));

            return 127;

        }

    };

    let pgid = i32::try_from(child.id()).expect("pid fits in pid_t");

    PGID.store(pgid, Ordering::Relaxed);

    for sig in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP] {

        unsafe { libc::signal(sig, forward as *const () as libc::sighandler_t) };

    }

    let mut out = grant.stream.try_clone().ok();

    if let Some(s) = &mut out {

        let _ = proto::send(s, &Request::Started { pid: pgid });

    }

    let killed = Arc::new(AtomicBool::new(false));

    if let Ok(mut reader) = Reader::new(&grant.stream) {

        let killed = killed.clone();

        thread::spawn(move || {

            while let Ok(Some(ev)) = reader.recv::<Event>() {

                if matches!(ev, Event::Kill) {

                    killed.store(true, Ordering::Relaxed);

                }

            }

        });

    }

    let started = Instant::now();

    let grace = need.time.mul_f64(cfg.overrun.time_grace);

    let loud_ram = Bytes((need.ram.0 as f64 * cfg.overrun.ram_slack) as u64);

    let mut peak = Bytes::ZERO;

    let mut last_sample: Option<Instant> = None;

    let mut term_sent: Option<Instant> = None;

    let mut timed_out = false;

    let mut warned_ram = false;

    let mut warned_time = false;

    let status = loop {

        match child.try_wait() {

            Ok(Some(status)) => break status,

            Ok(None) => {}

            Err(e) => {

                say(&format!("lost track of the child: {e}"));

                return 1;

            }

        }

        let elapsed = started.elapsed();

        if term_sent.is_none() {

            if killed.load(Ordering::Relaxed) {

                say("killed on request");

                term_sent = Some(Instant::now());

            } else if a.deadline.is_some_and(|d| elapsed > d.0) {

                say("deadline hit, killing it");

                timed_out = true;

                term_sent = Some(Instant::now());

            }

            if term_sent.is_some() {

                unsafe { libc::kill(-pgid, libc::SIGTERM) };

            }

        } else if term_sent.is_some_and(|t| t.elapsed() > Duration::from_secs(10)) {

            unsafe { libc::kill(-pgid, libc::SIGKILL) };

        }

        if last_sample.is_none_or(|t| t.elapsed() >= Duration::from_secs(2)) {

            last_sample = Some(Instant::now());

            let rss = sys::group_rss(pgid);

            peak = peak.max(rss);

            if let Some(s) = &mut out {

                let _ = proto::send(s, &Request::Usage { rss });

            }

            if !warned_ram && rss > loud_ram {

                warned_ram = true;

                say(&format!("using {rss} against a {} estimate", need.ram));

            }

        }

        if !warned_time && elapsed > grace {

            warned_time = true;

            say(&format!(

                "{} past its {} estimate and still going",

                Dur(elapsed),

                Dur(need.time)

            ));

        }

        thread::sleep(Duration::from_millis(250));

    };

    let wall = started.elapsed();

    if let Some(s) = &mut out {

        let _ = proto::send(s, &Request::Done);

    }

    history::record(

        history::key(&root, &a.cmd),

        peak,

        wall,

        cpu_time(),

        grant.cores,

    );

    artifacts::register(cfg, &root, true);

    if peak > loud_ram || wall > grace {

        say(&format!(

            "took {} and peaked at {}, asked for {} and {}. ask for more next time",

            Dur(wall),

            peak,

            Dur(need.time),

            need.ram

        ));

    }

    if timed_out {

        EXIT_TIMEOUT

    } else {

        code(status)

    }

}