//! init -- Foundry OS's first program (pid 1) under boot contract v2.
//!
//! When the kernel finds our ext2 root partition, it mounts it at `/` and
//! starts `/sbin/init` (this program) as process 1, with argv `["init"]`
//! (`["init", "selftest"]` in a headless build) and the environment
//! PATH=/sbin:/bin:/usr/bin:/mnt/disk0/bin, HOME=/, TERM=linux (plus
//! RUSTBELT_HEADLESS=1 when headless). See docs/boot-contract.md.
//!
//! Its jobs:
//!
//! 1. **Start the login shells.** rbsh (`/bin/rbsh`, the shell with drive
//!    letters and DOS commands) on the console, and another one on the
//!    second console `/dev/hvc0` if QEMU has one.
//! 2. **Respawn them.** Type `exit` in a shell and init starts a new one.
//!    (The kernel never restarts pid 1 itself, so init must never exit in
//!    normal use: when pid 1 exits, the kernel ends the QEMU run.)
//! 3. **Reap orphans.** When a process exits before its children, the
//!    kernel hands those children to pid 1 -- but only if pid 1 is called
//!    `init`, which is why this program is named init. A finished child
//!    stays a "zombie" until its parent collects its exit status with
//!    wait(); init's main loop waits for *any* child, so adopted orphans are
//!    collected too.
//! 4. **Headless mode** (`init selftest`, or RUSTBELT_HEADLESS=1), used by
//!    `make test`: run `rbsh /scripts/selftest.sh` on the console, wait for
//!    it and for every orphan it left behind, report how many memory frames
//!    are in use before and after (they must be equal: no leaks), and exit
//!    with the self-test's exit code. The kernel turns pid 1's exit(n) into
//!    QEMU exit code 33 + 2n, so exit(0) = 33 = "pass".
//!
//! init ignores the keyboard and termination signals (Ctrl+C on the console
//! must not kill pid 1); its children get the default handlers back.
#![no_std]
#![no_main]

extern crate alloc;

extern crate ulib;
ulib::entry!(main);

use ulib::*;

const SHELL: &str = "/bin/rbsh";
const SELFTEST: &str = "/scripts/selftest.sh";
const HVC0: &str = "/dev/hvc0";
/// Signals init ignores (and its children get back as SIG_DFL).
const IGNORED: [i32; 7] = [SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGTSTP, SIGTTIN, SIGTTOU];
const CLOCK_MONOTONIC: u64 = 1;

fn now_ms() -> u64 {
    let (s, ns) = clock_gettime(CLOCK_MONOTONIC);
    s * 1000 + ns / 1_000_000
}

/// Start `argv` (argv[0] is the path) in a child process. With `tty`, the
/// child gets a session of its own with that terminal as stdin, stdout,
/// stderr and controlling terminal (what getty does); without, it shares
/// init's console. Returns the child's pid, or -errno.
fn start(path: &str, argv: &[&str], tty: Option<&str>) -> i64 {
    let pid = fork();
    if pid != 0 {
        return pid; // init (or -errno if fork failed)
    }
    // The child.
    for sig in IGNORED {
        sigaction_raw(sig, SIG_DFL, 0, 0);
    }
    if let Some(dev) = tty {
        setsid();
        let fd = open(dev, O_RDWR); // becomes the controlling terminal
        if fd < 0 {
            eprintln!("[init] {}: {}", dev, strerror(fd));
            exit(126);
        }
        for target in 0..3 {
            dup2(fd, target);
        }
        if fd > 2 {
            close(fd);
        }
        println!("\nFoundry OS: shell on {} (type `exit` for a new one)", dev);
    }
    let err = execve(path, argv, &environ());
    eprintln!("[init] {}: {}", path, strerror(err));
    exit(127)
}

/// Print the mount table on one line, e.g. "/ ext2, /tmp ramfs, ...".
fn show_mounts() {
    let mut buf = [0u8; 1024];
    let n = mounts(&mut buf);
    if n > 0 {
        let text = core::str::from_utf8(&buf[..n as usize]).unwrap_or("");
        let list: alloc::vec::Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
        println!("[init] mounts: {}", list.join(", "));
    }
}

fn exists(path: &str) -> bool {
    let mut st = Stat::default();
    stat(path, &mut st) == 0
}

/// Headless mode: run the self-test and return its exit code.
fn selftest() -> i64 {
    // Settle first, then count: the frames in use now are the baseline.
    let before = sysinfo();
    let shell = start(SHELL, &["rbsh", SELFTEST], None);
    if shell < 0 {
        eprintln!("[init] can't start {}: {}", SHELL, strerror(shell));
        return 1;
    }
    let mut status = -1;
    let mut orphans = 0u64;
    // Wait for the shell, collecting any adopted orphan that ends first.
    loop {
        match wait4(-1, 0) {
            Ok((pid, st)) if pid == shell => {
                status = exit_code(st);
                break;
            }
            Ok(_) => orphans += 1,
            Err(e) if e == -EINTR => continue,
            Err(e) => {
                eprintln!("[init] wait: {}", strerror(e));
                break;
            }
        }
    }
    // Orphans may still be running (programs the self-test left in the
    // background whose parent has gone). Collect them for up to 5 seconds,
    // until only init itself is left.
    let deadline = now_ms() + 5000;
    loop {
        match wait4(-1, WNOHANG) {
            Ok((0, _)) => {}
            Ok(_) => {
                orphans += 1;
                continue;
            }
            Err(e) if e == -EINTR => continue,
            Err(_) => {} // ECHILD: no children left
        }
        let info = sysinfo();
        if (info.processes <= before.processes && info.threads <= before.threads) || now_ms() > deadline {
            break;
        }
        sleep_ms(10);
    }
    let after = sysinfo();
    if after.processes > before.processes {
        println!("[init] {} process(es) still running after 5 s", after.processes - before.processes);
    }
    println!("[init] self-test finished: exit code {}; init reaped {} adopted orphan(s)", status, orphans);
    println!(
        "[init] memory: {} frames in use before the self-test, {} after it and all its children exited",
        before.frames_in_use, after.frames_in_use
    );
    status
}

/// One terminal with a login shell on it.
struct Seat {
    name: &'static str,
    tty: Option<&'static str>,
    pid: i64,
    quick_deaths: u32,
    started_ms: u64,
}

impl Seat {
    fn spawn(&mut self) {
        self.started_ms = now_ms();
        self.pid = start(SHELL, &["rbsh"], self.tty);
        if self.pid < 0 {
            eprintln!("[init] can't start a shell on {}: {}", self.name, strerror(self.pid));
        }
    }
}

/// Normal mode: keep a shell on every terminal, forever.
fn supervise() -> ! {
    let mut seats = alloc::vec![Seat { name: "console", tty: None, pid: -1, quick_deaths: 0, started_ms: 0 }];
    if exists(HVC0) {
        seats.push(Seat { name: "hvc0", tty: Some(HVC0), pid: -1, quick_deaths: 0, started_ms: 0 });
    }
    for seat in seats.iter_mut() {
        seat.spawn();
    }
    loop {
        let (pid, st) = match wait4(-1, 0) {
            Ok(r) => r,
            Err(e) if e == -EINTR => continue,
            Err(_) => {
                sleep_ms(1000); // no children at all (every shell gave up)
                continue;
            }
        };
        // Not one of our shells: an adopted orphan, now reaped. Nothing else to do.
        let Some(seat) = seats.iter_mut().find(|s| s.pid == pid) else { continue };
        let code = exit_code(st);
        if now_ms() - seat.started_ms < 1000 {
            seat.quick_deaths += 1;
        } else {
            seat.quick_deaths = 0;
        }
        if seat.quick_deaths >= 5 || code == 126 || code == 127 {
            println!("[init] the shell on {} keeps failing (exit code {}); not restarting it", seat.name, code);
            seat.pid = -1;
            continue;
        }
        println!("[init] shell on {} exited (code {}); starting a new one", seat.name, code);
        if seat.quick_deaths > 0 {
            sleep_ms(500); // don't spin if it dies at once
        }
        seat.spawn();
    }
}

fn main() -> i64 {
    let headless = arg(1) == Some("selftest") || getenv("RUSTBELT_HEADLESS") == Some("1");
    for sig in IGNORED {
        sigaction_raw(sig, SIG_IGN, 0, 0);
    }
    let mut buf = [0u8; 6 * 65];
    let [_, _, release, _, _, _] = uname(&mut buf);
    println!(
        "[init] Foundry OS init, pid {}, on Rustbelt {}{}",
        getpid(),
        release,
        if headless { " (headless: running the self-test)" } else { "" }
    );
    if getpid() != 1 {
        eprintln!("[init] warning: I am not pid 1, so orphans won't be handed to me");
    }
    show_mounts();
    if headless {
        return selftest();
    }
    let mut motd = [0u8; 512];
    let fd = open("/etc/motd", 0);
    if fd >= 0 {
        let n = read(fd, &mut motd);
        close(fd);
        if n > 0 {
            write(1, &motd[..n as usize]);
        }
    }
    supervise()
}