//! sh -- the Rustbelt shell, a normal user program (/bin/sh on the disk).
//!
//!   sh               interactive: prompt, line editing, history
//!   sh script.sh     run the commands in a file (echoing each one)
//!
//! It reads a line, parses it (shparse.rs: words, `|`, `<`, `>`, `>>`, `&`,
//! drive letters), then either runs a *builtin* (cd, pwd, echo, ... -- these
//! must change the shell itself, or are just tiny) or *spawns* a program
//! from a `bin` directory, connecting pipes and files to its stdin/stdout.
//! Everything goes through the same syscalls any program could use.

#![no_std]
#![no_main]

extern crate alloc;

#[macro_use]
#[path = "ulib.rs"]
mod ulib;
#[path = "shparse.rs"]
mod shparse;

use alloc::{
    format,
    string::{String, ToString},
    vec::Vec,
};
use core::fmt::Write;
use shparse::Drive;
use ulib::*;

/// Built-in aliases (DOS names for the Unix-y programs).
/// Plus Linux names for the network tools.
const DEFAULT_ALIASES: &[(&str, &str)] = &[
    ("dir", "ls"), ("type", "cat"), ("del", "rm"), ("md", "mkdir"), ("copy", "cp"),
    ("ip", "ifconfig"), ("nslookup", "host"), ("wget", "fetch"),
];

const HELP: &str = "\
Rustbelt shell commands
  builtins   cd [dir]  pwd  echo [text]  ver  drives  alias [name=value]
             history  jobs  wait [pid]  cls / clear  help  exit [code]
  programs   ls [-l]  cat  wc  grep [-i] [-v]  mkdir  rm  touch  cp  sleep  sysinfo
  network    ifconfig (ip)  ping [-c n] host  host name (nslookup)
             fetch [-o file] http://... (wget)  httpd [-p port] [-n count]
  devices    random [-n bytes | -d sides [n]]  getty [-1] /dev/hvc0  tile [n] | tile -r
  (bin dirs) hello  demo  evil  spin [rounds]  overflow  child  sh
  run [--trace] prog args   run a program (--trace prints its syscalls)
  start prog args           run it in the background (or: prog args &)
  a | b | c                 pipes;  > file  >> file  < file  redirection
  A: is /, B: C: ... are the disks (see `drives`): cd B:, cat B:/readme.txt
";

/// A background job started with `start` or `&`.
struct Job {
    pid: i64,
    command: String,
}

/// One command of a pipeline, after alias, drive letter and `run`/`start`
/// processing.
struct Stage {
    argv: Vec<String>,
    stdin: Option<String>,
    stdout: Option<(String, bool)>, // (file, append)
    trace: bool,
    program_only: bool, // `run x` always runs a program, never a builtin
}

struct Shell {
    drives: Vec<Drive>,
    aliases: Vec<(String, String)>,
    history: Vec<String>,
    jobs: Vec<Job>,
    status: i64,            // exit code of the last command
    exit: Option<i64>,      // set by the `exit` builtin
}

fn main() -> i64 {
    let mut sh = Shell {
        drives: Vec::new(),
        aliases: DEFAULT_ALIASES.iter().map(|(a, b)| (a.to_string(), b.to_string())).collect(),
        history: Vec::new(),
        jobs: Vec::new(),
        status: 0,
        exit: None,
    };

    // Where do commands come from? A script file, or stdin. If stdin is the
    // console we're interactive; otherwise (a pipe, `sh < file`) we read
    // lines just like from a script.
    let fd = match arg(1) {
        Some(path) => {
            let fd = open(path, 0);
            if fd < 0 {
                eprintln!("sh: {}: {}", path, strerror(fd));
                return 127;
            }
            fd
        }
        None => 0,
    };
    let mut st = Stat::default();
    let interactive = fd == 0 && fstat(0, &mut st) == 0 && st.kind == KIND_CONSOLE;
    let mut input = LineReader { fd, buf: Vec::new(), eof: false };

    sh.refresh_drives();
    if interactive {
        let disks: Vec<String> = sh.drives.iter().map(|d| format!("{}: = {}", d.letter, d.path)).collect();
        println!("Rustbelt shell. Type `help` for commands. Drives: {}", disks.join(", "));
    }

    while sh.exit.is_none() {
        sh.refresh_drives();
        sh.reap_jobs();
        let prompt = sh.prompt();
        let line = if interactive {
            sh.edit_line(&prompt)
        } else {
            match input.read_line() {
                Some(line) => {
                    println!("{}{}", prompt, line); // show what's being run
                    line
                }
                None => break, // end of the script
            }
        };
        sh.run_line(&line, interactive);
    }
    sh.exit.unwrap_or(sh.status)
}

impl Shell {
    /// Ask the kernel for the mount table and (re)assign drive letters.
    fn refresh_drives(&mut self) {
        let mut buf = [0u8; 1024];
        let n = mounts(&mut buf).max(0) as usize;
        self.drives = shparse::assign_drives(core::str::from_utf8(&buf[..n]).unwrap_or(""));
    }

    fn cwd(&self) -> String {
        let mut buf = [0u8; 256];
        let n = getcwd(&mut buf).max(0) as usize;
        String::from_utf8_lossy(&buf[..n]).into_owned()
    }

    /// `A:/ $ `, `B:/notes $ `
    fn prompt(&self) -> String {
        format!("{} $ ", shparse::display(&self.cwd(), &self.drives))
    }

    /// Print "[done]" for background jobs that have finished (without
    /// blocking: WNOHANG returns -EAGAIN while a job is still running).
    fn reap_jobs(&mut self) {
        self.jobs.retain(|job| {
            let code = try_wait(job.pid);
            if code == -EAGAIN {
                return true;
            }
            println!("[done] pid {} ({}), exit code {}", job.pid, job.command, code);
            false
        });
    }

    /// Read a line from the keyboard with editing. The console is put in
    /// raw mode so we get every key (and arrow keys as ESC [ A..D) and do
    /// the echo ourselves; back to normal ("cooked") before running
    /// anything, so programs see the usual line-at-a-time input.
    fn edit_line(&mut self, prompt: &str) -> String {
        print!("{}", prompt);
        set_raw(0, true);
        let mut line: Vec<u8> = Vec::new();
        let mut hist = self.history.len(); // position while browsing history
        let mut esc = 0; // 0: normal, 1: after ESC, 2: after ESC [
        let mut buf = [0u8; 16];
        'read: loop {
            let n = read(0, &mut buf);
            if n <= 0 {
                break;
            }
            for &b in &buf[..n as usize] {
                match (esc, b) {
                    (0, 0x1b) => esc = 1,
                    (1, b'[') => esc = 2,
                    (1, _) => esc = 0,
                    (2, b'0'..=b'9' | b';') => {} // parameters, e.g. ESC [ 3 ~
                    (2, key) => {
                        esc = 0;
                        // Up / Down: replace the line with an older/newer one.
                        let new = match key {
                            b'A' if hist > 0 => hist - 1,
                            b'B' if hist < self.history.len() => hist + 1,
                            _ => continue, // Left/Right/Home/... not supported
                        };
                        hist = new;
                        let text = self.history.get(hist).map(|s| s.as_bytes()).unwrap_or(b"");
                        erase(line.len());
                        line = text.to_vec();
                        write(1, &line);
                    }
                    (_, b'\n' | b'\r') => {
                        println!();
                        break 'read;
                    }
                    (_, 0x08 | 0x7f) => {
                        if line.pop().is_some() {
                            erase(1);
                        }
                    }
                    (_, 0x20..=0x7e) => {
                        line.push(b);
                        write(1, &[b]); // echo
                    }
                    _ => {} // Tab, other control keys: ignored
                }
            }
        }
        set_raw(0, false);
        String::from_utf8_lossy(&line).into_owned()
    }

    fn run_line(&mut self, line: &str, interactive: bool) {
        let line = line.trim();
        if line.is_empty() {
            return;
        }
        if interactive && self.history.last().map(|h| h.as_str()) != Some(line) {
            self.history.push(line.to_string());
        }
        match shparse::parse(line) {
            Ok(Some(pipeline)) => self.status = self.run_pipeline(pipeline, line),
            Ok(None) => {} // just a comment
            Err(e) => {
                eprintln!("sh: {}", e);
                self.status = 2;
            }
        }
    }

    /// Words -> argv strings: drive letters translated (unless quoted),
    /// aliases expanded, `run [--trace]` / `start` prefixes removed.
    fn prepare(&self, cmd: &shparse::Command, background: &mut bool) -> Result<Stage, String> {
        let path = |w: &shparse::Word| match w.quoted {
            false => shparse::translate(&w.text, &self.drives).unwrap_or_else(|| w.text.clone()),
            true => w.text.clone(),
        };
        let mut argv: Vec<String> = cmd.words.iter().map(path).collect();
        let mut stage = Stage {
            argv: Vec::new(),
            stdin: cmd.stdin.as_ref().map(path),
            stdout: cmd.stdout.as_ref().map(|(w, append)| (path(w), *append)),
            trace: false,
            program_only: false,
        };
        // `run` / `start` prefixes (`start run --trace x` works too).
        loop {
            match argv.first().map(|s| s.as_str()) {
                Some("run") => {
                    argv.remove(0);
                    stage.program_only = true;
                    if argv.first().map(|s| s.as_str()) == Some("--trace") {
                        argv.remove(0);
                        stage.trace = true;
                    }
                }
                Some("start") => {
                    argv.remove(0);
                    stage.program_only = true;
                    *background = true;
                }
                _ => break,
            }
        }
        if argv.is_empty() {
            return Err(String::from("run/start: which program?"));
        }
        // Aliases: replace the command name (a few rounds, for aliases of
        // aliases; the limit stops `alias a=a` from looping forever).
        for _ in 0..5 {
            let Some((_, value)) = self.aliases.iter().find(|(name, _)| *name == argv[0]) else { break };
            let mut words: Vec<String> = value.split_whitespace().map(|s| s.to_string()).collect();
            if words.is_empty() || words[0] == argv[0] {
                break;
            }
            words.extend(argv.drain(1..));
            argv = words;
        }
        stage.argv = argv;
        Ok(stage)
    }

    /// Run `a < in | b | c > out`: create the pipes, spawn every program
    /// with the right stdin/stdout, then wait for them (unless background).
    /// Returns the exit code of the last command.
    fn run_pipeline(&mut self, pipeline: shparse::Pipeline, line: &str) -> i64 {
        let mut background = pipeline.background;
        let mut stages = Vec::new();
        for cmd in &pipeline.commands {
            match self.prepare(cmd, &mut background) {
                Ok(stage) => stages.push(stage),
                Err(e) => {
                    eprintln!("sh: {}", e);
                    return 2;
                }
            }
        }

        let mut opened: Vec<i64> = Vec::new(); // fds to close when done
        let mut children: Vec<(i64, String)> = Vec::new();
        let mut status = 0;
        let mut prev_pipe = -1; // read end of the pipe from the previous stage
        let mut last_pid = None; // the last stage, if it's a program
        for (i, stage) in stages.iter().enumerate() {
            // stdin: `< file`, else the previous pipe, else ours (-1).
            let mut stdin = prev_pipe;
            if let Some(file) = &stage.stdin {
                stdin = open(file, 0);
                if stdin < 0 {
                    eprintln!("sh: {}: {}", file, strerror(stdin));
                    status = 1;
                    break;
                }
                opened.push(stdin);
            }
            // stdout: `> file` / `>> file`, else a new pipe to the next
            // stage, else ours (-1).
            let mut stdout = -1;
            prev_pipe = -1;
            if i + 1 < stages.len() {
                let (r, w) = pipe().unwrap_or((-1, -1));
                opened.extend([r, w]);
                stdout = w;
                prev_pipe = r;
            }
            if let Some((file, append)) = &stage.stdout {
                let flags = O_CREAT | if *append { O_APPEND } else { O_TRUNC };
                stdout = open(file, flags);
                if stdout < 0 {
                    eprintln!("sh: {}: {}", file, strerror(stdout));
                    status = 1;
                    break;
                }
                opened.push(stdout);
            }

            if !stage.program_only {
                if let Some(code) = self.builtin(&stage.argv, if stdout < 0 { 1 } else { stdout }) {
                    status = code;
                    continue;
                }
            }
            match self.spawn(stage, stdin, stdout) {
                Ok(pid) => {
                    children.push((pid, stage.argv[0].clone()));
                    if i + 1 == stages.len() {
                        last_pid = Some(pid);
                    }
                }
                Err(code) => {
                    status = code;
                    break;
                }
            }
        }
        // Close our copies of the pipe ends and files: a reader only sees
        // end-of-file once *every* write end is closed, including ours.
        for fd in opened {
            close(fd);
        }

        if background {
            for (pid, name) in children {
                println!("[pid {}] {} started in the background", pid, name);
                self.jobs.push(Job { pid, command: line.to_string() });
            }
            return 0;
        }
        for (pid, _) in &children {
            let code = wait(*pid);
            if Some(*pid) == last_pid {
                status = code; // the last stage's exit code is the pipeline's
            }
        }
        if status != 0 && last_pid.is_some() {
            println!("[exit code {}]", status);
        }
        status
    }

    /// Find a program: names with a `/` are paths; others are looked up in
    /// the `bin` directory of every drive (our PATH): A:/bin, B:/bin, ...
    fn find_program(&self, name: &str) -> Option<String> {
        if name.contains('/') {
            return Some(name.to_string());
        }
        let mut st = Stat::default();
        self.drives.iter().map(|d| format!("{}/bin/{}", d.path.trim_end_matches('/'), name)).find(|p| stat(p, &mut st) == 0 && st.kind == KIND_FILE)
    }

    fn spawn(&self, stage: &Stage, stdin: i64, stdout: i64) -> Result<i64, i64> {
        let name = &stage.argv[0];
        let Some(path) = self.find_program(name) else {
            eprintln!("sh: {}: command not found (try `help`)", name);
            return Err(127);
        };
        // spawn takes the arguments as one NUL-separated string.
        let argv = stage.argv.join("\0");
        let opts = SpawnOptions { stdin, stdout, stderr: -1, flags: if stage.trace { SPAWN_TRACE } else { 0 } };
        let pid = spawn(&path, &argv, &opts);
        if pid < 0 {
            eprintln!("sh: {}: {}", name, strerror(pid));
            return Err(126);
        }
        Ok(pid)
    }

    /// Run a builtin, writing its output to `out`. `None` if `argv[0]`
    /// isn't one.
    fn builtin(&mut self, argv: &[String], out: i64) -> Option<i64> {
        let mut out = Fd(out);
        let arg1 = argv.get(1).map(|s| s.as_str());
        let code = match argv[0].as_str() {
            "cd" => {
                let dir = arg1.unwrap_or("/");
                let r = chdir(dir);
                if r < 0 {
                    match shparse::drive_letter(dir) {
                        // Still has its letter: translate() didn't know it.
                        Some(letter) => eprintln!("cd: there is no drive {}: (see `drives`)", letter),
                        None => eprintln!("cd: {}: {}", dir, strerror(r)),
                    }
                    1
                } else {
                    0
                }
            }
            "pwd" => {
                let _ = writeln!(out, "{}", self.cwd());
                0
            }
            "echo" => {
                let _ = writeln!(out, "{}", argv[1..].join(" "));
                0
            }
            "help" => {
                let _ = write!(out, "{}", HELP);
                0
            }
            "exit" => {
                self.exit = Some(arg1.and_then(|s| s.parse().ok()).unwrap_or(self.status));
                0
            }
            "cls" | "clear" => {
                let _ = write!(out, "\x0c"); // form feed: the console clears the screen
                0
            }
            "ver" => {
                let mut buf = [0u8; 6 * 65];
                let [sysname, _node, release, version, machine, _] = uname(&mut buf);
                let _ = writeln!(out, "{} {} ({}) {}", sysname, release, version, machine);
                0
            }
            "drives" => {
                for d in &self.drives {
                    let _ = writeln!(out, "{}:  {:<12} {}", d.letter, d.path, d.fs);
                }
                0
            }
            "alias" => match arg1 {
                None => {
                    for (name, value) in &self.aliases {
                        let _ = writeln!(out, "{} = {}", name, value);
                    }
                    0
                }
                Some(def) => match def.split_once('=') {
                    Some((name, value)) => {
                        self.aliases.retain(|(n, _)| n != name);
                        self.aliases.push((name.to_string(), value.to_string()));
                        0
                    }
                    None => {
                        eprintln!("alias: use alias name=value");
                        2
                    }
                },
            },
            "history" => {
                for (i, line) in self.history.iter().enumerate() {
                    let _ = writeln!(out, "{:4}  {}", i + 1, line);
                }
                0
            }
            "jobs" => {
                self.reap_jobs();
                for job in &self.jobs {
                    let _ = writeln!(out, "[pid {}] running  {}", job.pid, job.command);
                }
                0
            }
            "wait" => {
                // `wait` waits for all background jobs, `wait 7` for pid 7.
                let only: Option<i64> = arg1.and_then(|s| s.parse().ok());
                let mut code = 0;
                self.jobs.retain(|job| {
                    if only.is_some_and(|pid| pid != job.pid) {
                        return true;
                    }
                    code = wait(job.pid);
                    let _ = writeln!(out, "[done] pid {} ({}), exit code {}", job.pid, job.command, code);
                    false
                });
                code
            }
            _ => return None,
        };
        Some(code)
    }
}

/// Erase `n` characters before the cursor (the console treats \x08 as
/// "back one column and blank it").
fn erase(n: usize) {
    for _ in 0..n {
        write(1, b"\x08");
    }
}

/// Reads a file (or pipe) line by line.
struct LineReader {
    fd: i64,
    buf: Vec<u8>,
    eof: bool,
}

impl LineReader {
    fn read_line(&mut self) -> Option<String> {
        loop {
            if let Some(i) = self.buf.iter().position(|&b| b == b'\n') {
                let line: Vec<u8> = self.buf.drain(..=i).collect();
                return Some(String::from_utf8_lossy(&line[..i]).trim_end_matches('\r').to_string());
            }
            if self.eof {
                // The last line may lack a final newline.
                if self.buf.is_empty() {
                    return None;
                }
                let line = String::from_utf8_lossy(&self.buf).into_owned();
                self.buf.clear();
                return Some(line);
            }
            let mut chunk = [0u8; 512];
            let n = read(self.fd, &mut chunk);
            if n <= 0 {
                self.eof = true;
            } else {
                self.buf.extend_from_slice(&chunk[..n as usize]);
            }
        }
    }
}