//! ruststd - Phase 2 Rust smoke test for MontaukOS: the native `std`.
//!
//! Runs a set of PASS/FAIL checks over the std platform layer: collections and
//! formatting, args and environment, the filesystem, threads (joined, scoped
//! and detached), the futex-backed sync primitives, thread_local! destructors,
//! time, and child processes (exit codes, arguments, environment, working
//! directory, std::process::exit from a worker thread, panics).
//!
//! `ruststd net <host> [port]` additionally fetches `http://<host>/` over TCP.
//! The `child-*` arguments are the child-process halves of the checks.

use std::collections::{BTreeMap, HashMap};
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::process::{Command, ExitCode};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Barrier, Condvar, Mutex, OnceLock, RwLock, mpsc};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use std::{env, thread};

static FAILURES: AtomicUsize = AtomicUsize::new(0);

fn check(ok: bool, name: &str) {
    println!("{}  {}", if ok { "PASS" } else { "FAIL" }, name);
    if !ok {
        FAILURES.fetch_add(1, Ordering::Relaxed);
    }
}

fn main() -> ExitCode {
    let args: Vec<String> = env::args().collect();
    match args.get(1).map(String::as_str) {
        Some(mode) if mode.starts_with("child-") => return ExitCode::from(child(mode, &args)),
        Some("net") => return net(&args),
        _ => {}
    }

    println!("ruststd: Rust std on MontaukOS");
    println!("argv = {args:?}");

    basics();
    environment();
    filesystem();
    threads();
    sync();
    thread_locals();
    time();
    processes(&args);

    let failures = FAILURES.load(Ordering::Relaxed);
    if failures == 0 {
        println!("ruststd: all passed");
        ExitCode::SUCCESS
    } else {
        println!("ruststd: {failures} FAILED");
        ExitCode::FAILURE
    }
}

// ====================================================================
// Child-process halves
// ====================================================================

fn child(mode: &str, args: &[String]) -> u8 {
    match mode {
        "child-exit" => 7,
        "child-args" => {
            let ok = args[2..] == ["with space", "", "it's", "say \"hi\"", "plain"];
            if ok { 8 } else { 1 }
        }
        "child-env" => {
            if env::var("RUSTSTD_VAR").as_deref() == Ok("hello there") { 9 } else { 1 }
        }
        "child-cwd" => {
            let cwd = env::current_dir().unwrap_or_default();
            if cwd.ends_with("os") { 10 } else { 1 }
        }
        "child-exit-thread" => {
            // std::process::exit from a worker must end the whole process.
            let h = thread::spawn(|| std::process::exit(11));
            let _ = h.join();
            thread::sleep(Duration::from_secs(5));
            1
        }
        "child-panic-thread" => {
            let h = thread::spawn(|| panic!("requested panic on a worker thread"));
            let _ = h.join();
            thread::sleep(Duration::from_secs(5));
            1
        }
        _ => 2,
    }
}

// ====================================================================
// Checks
// ====================================================================

fn basics() {
    let mut map = HashMap::new();
    for (i, w) in "Montauk speaks Rust with std".split(' ').enumerate() {
        map.insert(w.to_string(), i);
    }
    check(map["Rust"] == 2 && map.len() == 5, "HashMap (RandomState keys from SYS_GETRANDOM)");

    let sorted: BTreeMap<_, _> = map.iter().map(|(k, v)| (*v, k.clone())).collect();
    check(sorted.values().next().map(String::as_str) == Some("Montauk"), "BTreeMap");

    check(
        format!("{:>8.3}|{:e}|{:#06x}", std::f64::consts::PI, 1234.5f64, 255) == "   3.142|1.2345e3|0x00ff",
        "float and integer formatting",
    );

    let mut out = std::io::stdout().lock();
    check(writeln!(out, "  (stdout lock + writeln!)").is_ok(), "io::stdout().lock()");
    eprintln!("  (eprintln! goes to the terminal too)");
}

fn environment() {
    let exe = env::current_exe();
    check(exe.as_ref().is_ok_and(|p| p.is_absolute()), "env::current_exe() is absolute");
    println!("  current_exe = {exe:?}");

    let cwd = env::current_dir();
    check(cwd.is_ok(), "env::current_dir()");

    unsafe { env::set_var("RUSTSTD_TEST", "value=with=equals") };
    check(env::var("RUSTSTD_TEST").as_deref() == Ok("value=with=equals"), "set_var / var");
    check(env::vars().any(|(k, v)| k == "RUSTSTD_TEST" && v == "value=with=equals"), "env::vars() lists it");
    unsafe { env::remove_var("RUSTSTD_TEST") };
    check(env::var_os("RUSTSTD_TEST").is_none(), "remove_var");

    let joined = env::join_paths(["0:/os", "0:/games"]).unwrap();
    let split: Vec<PathBuf> = env::split_paths(&joined).collect();
    check(joined == "0:/os;0:/games" && split.len() == 2, "join_paths / split_paths use ';'");

    check(Path::new("0:/os/shell.elf").is_absolute(), "drive paths are absolute");
    check(Path::new("0:/os").join("0:/apps") == Path::new("0:/apps"), "joining a drive path replaces");
    check(Path::new("0:/os/shell.elf").parent() == Some(Path::new("0:/os")), "Path::parent");
}

fn filesystem() {
    let base = env::temp_dir().join(format!("ruststd-{}", std::process::id()));
    let _ = fs::remove_dir_all(&base);
    check(fs::create_dir_all(base.join("sub/deeper")).is_ok(), "create_dir_all");
    check(fs::create_dir(&base).is_err_and(|e| e.kind() == std::io::ErrorKind::AlreadyExists),
          "create_dir on an existing dir fails with AlreadyExists");

    let file = base.join("hello.txt");
    check(fs::write(&file, b"line one\nline two\n").is_ok(), "fs::write");
    check(fs::read_to_string(&file).ok().as_deref() == Some("line one\nline two\n"), "fs::read_to_string");

    let appended = OpenOptions::new().append(true).open(&file).and_then(|mut f| f.write_all(b"line three\n"));
    check(appended.is_ok(), "append");
    let lines: Vec<String> = BufReader::new(File::open(&file).unwrap()).lines().map(Result::unwrap).collect();
    check(lines == ["line one", "line two", "line three"], "BufReader::lines after append");

    let meta = fs::metadata(&file);
    check(meta.as_ref().is_ok_and(|m| m.is_file() && m.len() == 29), "metadata: is_file, len");
    if let Ok(m) = &meta {
        println!("  modified = {:?}", m.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()));
    }
    check(fs::metadata(&base).is_ok_and(|m| m.is_dir()), "metadata: is_dir");

    // Seek + overwrite in place.
    {
        let mut f = OpenOptions::new().read(true).write(true).open(&file).unwrap();
        f.seek(SeekFrom::Start(5)).unwrap();
        f.write_all(b"ONE").unwrap();
        f.seek(SeekFrom::End(-6)).unwrap();
        let mut tail = String::new();
        f.read_to_string(&mut tail).unwrap();
        check(tail == "three\n", "seek from end + read");
    }
    check(fs::read_to_string(&file).unwrap().starts_with("line ONE\n"), "seek + overwrite");

    // Truncate on open.
    fs::write(&file, b"short").unwrap();
    check(fs::read(&file).ok().as_deref() == Some(&b"short"[..]), "File::create truncates");

    check(fs::copy(&file, base.join("sub/copy.txt")).ok() == Some(5), "fs::copy");
    check(fs::rename(base.join("sub/copy.txt"), base.join("sub/moved.txt")).is_ok(), "fs::rename");

    let mut names: Vec<String> = fs::read_dir(&base)
        .map(|rd| rd.filter_map(|e| e.ok()).map(|e| e.file_name().to_string_lossy().into_owned()).collect())
        .unwrap_or_default();
    names.sort();
    check(names == ["hello.txt", "sub"], "read_dir lists entries");
    let sub_types: Vec<(String, bool)> = fs::read_dir(base.join("sub"))
        .unwrap()
        .map(|e| e.unwrap())
        .map(|e| (e.file_name().to_string_lossy().into_owned(), e.file_type().unwrap().is_dir()))
        .collect();
    check(sub_types.contains(&("deeper".into(), true)) && sub_types.contains(&("moved.txt".into(), false)),
          "DirEntry::file_type");

    check(File::open(base.join("missing")).is_err_and(|e| e.kind() == std::io::ErrorKind::NotFound),
          "open missing file -> NotFound");
    check(fs::remove_file(&file).is_ok() && !file.exists(), "remove_file");
    check(fs::remove_dir_all(&base).is_ok() && !base.exists(), "remove_dir_all");
}

fn threads() {
    let handles: Vec<_> = (0..4).map(|i| thread::spawn(move || i * 10)).collect();
    let results: Vec<i32> = handles.into_iter().map(|h| h.join().unwrap()).collect();
    check(results == [0, 10, 20, 30], "spawn + join return values");

    let named = thread::Builder::new().name("worker".into()).spawn(|| {
        thread::current().name().map(str::to_owned)
    });
    check(named.unwrap().join().unwrap().as_deref() == Some("worker"), "named thread");

    let mut data = vec![1, 2, 3, 4];
    thread::scope(|s| {
        for x in data.iter_mut() {
            s.spawn(move || *x *= 2);
        }
    });
    check(data == [2, 4, 6, 8], "thread::scope borrows");

    // Detached threads hold a kernel slot until reaped; 300 of them would
    // exhaust the 256-slot table if they were not.
    let done = Arc::new(AtomicUsize::new(0));
    for _ in 0..15 {
        for _ in 0..20 {
            let done = done.clone();
            drop(thread::spawn(move || {
                done.fetch_add(1, Ordering::Relaxed);
            }));
        }
        thread::sleep(Duration::from_millis(20));
    }
    let deadline = Instant::now() + Duration::from_secs(5);
    while done.load(Ordering::Relaxed) < 300 && Instant::now() < deadline {
        thread::sleep(Duration::from_millis(10));
    }
    check(done.load(Ordering::Relaxed) == 300, "300 detached threads are reaped");

    let big = thread::Builder::new().stack_size(4 << 20).spawn(|| {
        let buf = [7u8; 2 << 20];
        buf.iter().map(|&b| b as usize).sum::<usize>()
    });
    check(big.unwrap().join().unwrap() == 7 * (2 << 20), "4 MiB thread stack");
}

fn sync() {
    const THREADS: usize = 4;
    const ITERS: usize = 10_000;

    let counter = Arc::new(Mutex::new(0usize));
    let handles: Vec<_> = (0..THREADS)
        .map(|_| {
            let counter = counter.clone();
            thread::spawn(move || {
                for i in 0..ITERS {
                    let mut g = counter.lock().unwrap();
                    if i % 1000 == 0 {
                        thread::yield_now();
                    }
                    *g += 1;
                }
            })
        })
        .collect();
    handles.into_iter().for_each(|h| h.join().unwrap());
    check(*counter.lock().unwrap() == THREADS * ITERS, "Mutex keeps the counter exact");

    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();
    let waiter = thread::spawn(move || {
        let (lock, cvar) = &*pair2;
        let guard = cvar.wait_while(lock.lock().unwrap(), |ready| !*ready).unwrap();
        *guard
    });
    thread::sleep(Duration::from_millis(20));
    *pair.0.lock().unwrap() = true;
    pair.1.notify_all();
    check(waiter.join().unwrap(), "Condvar wait/notify");

    let (lock, cvar) = (Mutex::new(()), Condvar::new());
    let t0 = Instant::now();
    let (_g, res) = cvar.wait_timeout(lock.lock().unwrap(), Duration::from_millis(30)).unwrap();
    check(res.timed_out() && t0.elapsed() >= Duration::from_millis(30), "Condvar wait_timeout");

    let rw = Arc::new(RwLock::new(vec![1]));
    let readers: Vec<_> = (0..3)
        .map(|_| {
            let rw = rw.clone();
            thread::spawn(move || rw.read().unwrap().len())
        })
        .collect();
    rw.write().unwrap().push(2);
    let lens: Vec<usize> = readers.into_iter().map(|h| h.join().unwrap()).collect();
    check(lens.iter().all(|&l| l == 1 || l == 2) && rw.read().unwrap().len() == 2, "RwLock");

    static ONCE: OnceLock<usize> = OnceLock::new();
    let inits = Arc::new(AtomicUsize::new(0));
    let hs: Vec<_> = (0..4)
        .map(|_| {
            let inits = inits.clone();
            thread::spawn(move || *ONCE.get_or_init(|| inits.fetch_add(1, Ordering::SeqCst) + 42))
        })
        .collect();
    let vals: Vec<usize> = hs.into_iter().map(|h| h.join().unwrap()).collect();
    check(vals.iter().all(|&v| v == 42) && inits.load(Ordering::SeqCst) == 1, "OnceLock initializes once");

    let (tx, rx) = mpsc::channel();
    for t in 0..3 {
        let tx = tx.clone();
        thread::spawn(move || {
            for i in 0..100 {
                tx.send(t * 1000 + i).unwrap();
            }
        });
    }
    drop(tx);
    let received: Vec<i32> = rx.iter().collect();
    check(received.len() == 300, "mpsc channel with 3 senders");

    let barrier = Arc::new(Barrier::new(3));
    let hs: Vec<_> = (0..3)
        .map(|_| {
            let b = barrier.clone();
            thread::spawn(move || b.wait().is_leader())
        })
        .collect();
    let leaders = hs.into_iter().map(|h| h.join().unwrap()).filter(|&l| l).count();
    check(leaders == 1, "Barrier");

    let parked = thread::spawn(|| {
        thread::park();
        true
    });
    thread::sleep(Duration::from_millis(10));
    parked.thread().unpark();
    check(parked.join().unwrap(), "park / unpark");
}

fn thread_locals() {
    static DROPS: AtomicUsize = AtomicUsize::new(0);

    struct Tracker(u32);
    impl Drop for Tracker {
        fn drop(&mut self) {
            DROPS.fetch_add(1, Ordering::SeqCst);
        }
    }

    thread_local! {
        static COUNTER: std::cell::Cell<u32> = const { std::cell::Cell::new(5) };
        static TRACKER: Tracker = Tracker(1);
    }

    COUNTER.with(|c| c.set(100));
    let seen = thread::spawn(|| {
        TRACKER.with(|t| assert_eq!(t.0, 1));
        COUNTER.with(|c| c.get())
    })
    .join()
    .unwrap();
    check(seen == 5 && COUNTER.with(|c| c.get()) == 100, "thread_local! is per thread");
    check(DROPS.load(Ordering::SeqCst) == 1, "thread_local! destructor runs at thread exit");
}

fn time() {
    let t0 = Instant::now();
    thread::sleep(Duration::from_millis(50));
    let dt = t0.elapsed();
    check(dt >= Duration::from_millis(50) && dt < Duration::from_secs(2), "Instant + sleep(50 ms)");

    let now = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0);
    // 2025-01-01T00:00:00Z
    check(now > 1_735_689_600, "SystemTime::now() is after 2025");
    println!("  unix time = {now}");
}

fn processes(args: &[String]) {
    let Ok(me) = env::current_exe() else {
        check(false, "current_exe for child processes");
        return;
    };
    let _ = args;
    let run = |cmd: &mut Command| cmd.status().map(|s| s.code());

    check(run(Command::new(&me).arg("child-exit")).ok() == Some(Some(7)), "child exit code");
    check(
        run(Command::new(&me).args(["child-args", "with space", "", "it's", "say \"hi\"", "plain"])).ok() == Some(Some(8)),
        "child arguments with spaces and quotes",
    );
    check(
        run(Command::new(&me).arg("child-env").env("RUSTSTD_VAR", "hello there")).ok() == Some(Some(9)),
        "child environment",
    );
    check(
        run(Command::new(&me).arg("child-cwd").current_dir("0:/os")).ok() == Some(Some(10)),
        "child working directory",
    );
    check(
        run(Command::new(&me).arg("child-exit-thread")).ok() == Some(Some(11)),
        "process::exit from a worker thread ends the process",
    );
    let panicked = Command::new(&me).arg("child-panic-thread").status();
    check(
        panicked.as_ref().is_ok_and(|s| s.code().is_none() && s.to_string().contains("signal 6")),
        "panic on a worker thread aborts the process",
    );
    if let Ok(s) = panicked {
        println!("  panicking child: {s}");
    }

    let mut child = Command::new(&me).arg("child-exit").spawn().unwrap();
    let mut polled = None;
    for _ in 0..200 {
        if let Ok(Some(s)) = child.try_wait() {
            polled = Some(s);
            break;
        }
        thread::sleep(Duration::from_millis(10));
    }
    check(polled.and_then(|s| s.code()) == Some(7), "Child::try_wait");

    check(Command::new("no-such-program").status().is_err(), "spawning a missing program fails");
    check(
        Command::new(&me).arg("child-exit").output().is_err_and(|e| e.kind() == std::io::ErrorKind::Unsupported),
        "output() reports Unsupported (no stdio redirection yet)",
    );
}

// ====================================================================
// Network (optional: needs a link)
// ====================================================================

fn net(args: &[String]) -> ExitCode {
    use std::net::{TcpStream, ToSocketAddrs};

    let Some(host) = args.get(2) else {
        println!("usage: ruststd net <host> [port]");
        return ExitCode::FAILURE;
    };
    let port: u16 = args.get(3).and_then(|p| p.parse().ok()).unwrap_or(80);

    let addrs: Vec<_> = match (host.as_str(), port).to_socket_addrs() {
        Ok(a) => a.collect(),
        Err(e) => {
            println!("resolve {host}: {e}");
            return ExitCode::FAILURE;
        }
    };
    println!("{host} -> {addrs:?}");

    let mut stream = match TcpStream::connect((host.as_str(), port)) {
        Ok(s) => s,
        Err(e) => {
            println!("connect: {e}");
            return ExitCode::FAILURE;
        }
    };
    stream.set_read_timeout(Some(Duration::from_secs(10))).unwrap();
    let request = format!("GET / HTTP/1.0\r\nHost: {host}\r\nConnection: close\r\n\r\n");
    if let Err(e) = stream.write_all(request.as_bytes()) {
        println!("send: {e}");
        return ExitCode::FAILURE;
    }
    let mut response = Vec::new();
    match stream.read_to_end(&mut response) {
        Ok(n) => {
            let text = String::from_utf8_lossy(&response);
            println!("received {n} bytes; status line: {:?}", text.lines().next().unwrap_or(""));
            ExitCode::SUCCESS
        }
        Err(e) => {
            println!("recv: {e}");
            ExitCode::FAILURE
        }
    }
}