//! Processes = a thread + its own address space + a file descriptor table
//! (+ arguments, a current directory and a parent).
//!
//! Starting one (`spawn_image`):
//!   1. new level-4 page table sharing the kernel mappings (memory.rs)
//!   2. load the ELF program's segments into its user part (elf.rs)
//!   3. map a user stack just below `USER_END`
//!   4. fd table: 0, 1, 2 = whatever the parent chose (console by default,
//!      or e.g. a pipe -- that's how the shell will do `a | b` and `> file`)
//!   5. put it in the process table and spawn a thread with that CR3. The
//!      thread starts in the kernel (`user_thread_entry`), copies the
//!      arguments onto the user stack and "returns" into ring 3 with `iretq`.
//!
//! Why `spawn(path, argv)` instead of Unix `fork()` + `exec()`? fork copies
//! the whole address space (or needs copy-on-write page tables to be fast),
//! and then exec throws the copy away again. spawn ("create a new process
//! running this program") is one simple step, and it's what Windows
//! (CreateProcess) and POSIX's own `posix_spawn` offer. Redirection, which
//! Unix shells do between fork and exec, is done with spawn's stdio options.
//!
//! Exit and wait: a process's exit status is stored in it and its `exited`
//! wait queue is woken; the parent's `wait(pid)` blocks on that queue, then
//! removes the child from the process table ("reaps" it -- until then it's
//! what Unix calls a zombie).
//!
//! Not done yet (exercises): freeing a dead process's frames and page tables,
//! reparenting orphans, signals/kill, more than one thread per process.

use crate::elf::{self, ElfError};
use crate::errno::Errno;
use crate::fs::{self, file::File};
use crate::memory::{self, USER_END};
use crate::thread::{self, WaitQueue};
use crate::{serial_println, tty::Console};
use alloc::{boxed::Box, collections::BTreeMap, string::String, sync::Arc, vec, vec::Vec};
use core::sync::atomic::{AtomicU64, Ordering};
use spin::Mutex;
use x86_64::VirtAddr;
use x86_64::instructions::interrupts::without_interrupts;
use x86_64::structures::paging::{Page, PageTableFlags, Size4KiB};

/// The user stack: 16 KiB, ending one (unmapped) page below USER_END.
pub const USER_STACK_TOP: u64 = USER_END - 0x1000;
pub const USER_STACK_SIZE: u64 = 16 * 1024;
/// Max open files per process.
const MAX_FDS: usize = 16;
/// Max total size of a program's arguments (they go on its 16 KiB stack).
pub const MAX_ARGS_SIZE: usize = 2048;
/// Exit code `wait` reports for a process the kernel killed (like a shell
/// shows 128 + SIGSEGV(11) for a segfault).
pub const KILLED_EXIT_CODE: i64 = 139;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExitStatus {
    /// Called exit(code).
    Exited(i64),
    /// Killed by the kernel (e.g. "page fault").
    Killed(&'static str),
}

impl ExitStatus {
    /// The number `wait` returns.
    pub fn code(self) -> i64 {
        match self {
            ExitStatus::Exited(code) => code,
            ExitStatus::Killed(_) => KILLED_EXIT_CODE,
        }
    }
}

pub struct Process {
    pub pid: u64,
    /// Parent pid; 0 = started by the kernel itself.
    pub parent: u64,
    pub name: String,
    /// Print every syscall this process makes (`spawn` flag; `run --trace`).
    pub trace: bool,
    /// Current directory (always a normalized absolute path).
    pub cwd: Mutex<String>,
    entry: u64,
    args: Vec<String>,
    fds: Mutex<Vec<Option<Arc<dyn File>>>>,
    /// Locked with interrupts off only (it's read in a wait_until condition).
    status: Mutex<Option<ExitStatus>>,
    exited: WaitQueue,
}

static NEXT_PID: AtomicU64 = AtomicU64::new(1);
/// All processes that haven't been reaped yet, by pid.
static PROCESSES: Mutex<BTreeMap<u64, Arc<Process>>> = Mutex::new(BTreeMap::new());

/// Everything needed to start a process besides the program itself.
pub struct SpawnArgs {
    pub args: Vec<String>, // argv; args[0] is the program name by convention
    pub stdio: [Arc<dyn File>; 3],
    pub cwd: String,
    pub parent: u64,
    pub trace: bool,
}

impl SpawnArgs {
    /// Defaults for a process started by the kernel: console stdio, cwd "/".
    pub fn kernel(name: &str) -> SpawnArgs {
        let console: Arc<dyn File> = Arc::new(Console);
        SpawnArgs {
            args: vec![String::from(name)],
            stdio: [console.clone(), console.clone(), console],
            cwd: String::from("/"),
            parent: 0,
            trace: false,
        }
    }
}

/// Load the program at `path` (absolute) from the VFS and start it.
pub fn spawn_path(path: &str, spawn: SpawnArgs) -> Result<Arc<Process>, Errno> {
    let image = fs::read_all(path)?; // ENOENT, EISDIR, EIO, ...
    let name = path.rsplit('/').next().unwrap_or(path);
    spawn_image(name, &image, spawn).map_err(|e| match e {
        ElfError::OutOfMemory => Errno::ENOMEM,
        _ => Errno::ENOEXEC, // "Exec format error"
    })
}

/// Start ELF program `elf_image` from the kernel (console stdio, no args).
pub fn spawn(name: &str, elf_image: &[u8]) -> Result<Arc<Process>, ElfError> {
    spawn_image(name, elf_image, SpawnArgs::kernel(name))
}

/// Load `elf_image` into a new address space and start it.
pub fn spawn_image(name: &str, elf_image: &[u8], spawn: SpawnArgs) -> Result<Arc<Process>, ElfError> {
    let l4 = memory::new_address_space().ok_or(ElfError::OutOfMemory)?;
    let entry = elf::load(elf_image, l4)?;

    let mut stack_flags = PageTableFlags::WRITABLE;
    if memory::nx_enabled() {
        stack_flags |= PageTableFlags::NO_EXECUTE;
    }
    let first = Page::<Size4KiB>::containing_address(VirtAddr::new(USER_STACK_TOP - USER_STACK_SIZE));
    let last = Page::<Size4KiB>::containing_address(VirtAddr::new(USER_STACK_TOP - 1));
    for page in Page::range_inclusive(first, last) {
        memory::map_user_page(l4, page, stack_flags).ok_or(ElfError::OutOfMemory)?;
    }

    let [stdin, stdout, stderr] = spawn.stdio;
    let process = Arc::new(Process {
        pid: NEXT_PID.fetch_add(1, Ordering::Relaxed),
        parent: spawn.parent,
        name: String::from(name),
        trace: spawn.trace,
        cwd: Mutex::new(spawn.cwd),
        entry,
        args: spawn.args,
        fds: Mutex::new(vec![Some(stdin), Some(stdout), Some(stderr)]),
        status: Mutex::new(None),
        exited: WaitQueue::new(),
    });
    PROCESSES.lock().insert(process.pid, process.clone());
    serial_println!("[process] started pid {} ({}), entry {:#x}", process.pid, name, entry);
    // Thread names are &'static str; leaking a few bytes per process is fine here.
    let thread_name: &'static str = Box::leak(String::from(name).into_boxed_str());
    thread::spawn_with(thread_name, user_thread_entry, l4, Some(process.clone()));
    Ok(process)
}

/// First code of every process thread (still ring 0, but already running
/// with the process's CR3 and kernel stack).
fn user_thread_entry() {
    let process = thread::current_process().expect("process thread without process");
    // Copy the arguments onto the top of the user stack (we can write there
    // directly: this process's address space is active). Layout, top down:
    //   the argument strings' bytes (not NUL-terminated)
    //   argv: [ptr, len] pairs, one per argument   <- rsi
    //   (stack continues below)
    // and rdi = argc. Rust's `&str` is ptr+len, so pairs are handier than
    // C's NUL-terminated `char **argv`.
    let mut sp = USER_STACK_TOP;
    let mut pairs: Vec<[u64; 2]> = Vec::with_capacity(process.args.len());
    for arg in &process.args {
        sp -= arg.len() as u64;
        unsafe { core::ptr::copy_nonoverlapping(arg.as_ptr(), sp as *mut u8, arg.len()) };
        pairs.push([sp, arg.len() as u64]);
    }
    sp &= !0xf; // 16-byte align
    sp -= 16 * pairs.len() as u64;
    let argv = sp;
    unsafe { core::ptr::copy_nonoverlapping(pairs.as_ptr(), argv as *mut [u64; 2], pairs.len()) };
    // Functions expect rsp % 16 == 8 on entry (as if a return address had
    // just been pushed), so start 8 bytes below the 16-aligned spot.
    let (entry, argc) = (process.entry, pairs.len() as u64);
    drop(pairs);
    drop(process); // enter_user_mode never returns: don't leak the Arc
    unsafe { enter_user_mode(entry, sp - 8, argc, argv) }
}

/// Jump to ring 3. `iretq` pops RIP, CS, RFLAGS, RSP, SS -- we push a fake
/// interrupt frame with user selectors (RPL 3), so "returning" from it drops
/// the privilege level. RFLAGS = 0x202: interrupts enabled (+ the always-1 bit).
/// rdi/rsi carry `_start(argc, argv)`'s arguments.
unsafe fn enter_user_mode(rip: u64, rsp: u64, argc: u64, argv: u64) -> ! {
    let s = crate::gdt::selectors();
    unsafe {
        core::arch::asm!(
            "push {ss}",
            "push {rsp}",
            "push {rflags}",
            "push {cs}",
            "push {rip}",
            // Don't leak kernel values to user space in registers.
            "xor eax, eax", "xor ebx, ebx", "xor ecx, ecx", "xor edx, edx",
            "xor ebp, ebp",
            "xor r8d, r8d", "xor r9d, r9d", "xor r10d, r10d", "xor r11d, r11d",
            "xor r12d, r12d", "xor r13d, r13d", "xor r14d, r14d", "xor r15d, r15d",
            "iretq",
            ss = in(reg) s.user_data.0 as u64,
            rsp = in(reg) rsp,
            rflags = in(reg) 0x202u64,
            cs = in(reg) s.user_code.0 as u64,
            rip = in(reg) rip,
            in("rdi") argc,
            in("rsi") argv,
            options(noreturn)
        );
    }
}

impl Process {
    /// The open file behind `fd`.
    pub fn file(&self, fd: u64) -> Result<Arc<dyn File>, Errno> {
        let fds = self.fds.lock();
        fds.get(fd as usize).and_then(|f| f.clone()).ok_or(Errno::EBADF)
    }

    /// Put `file` into the lowest free fd slot (like Unix).
    pub fn add_file(&self, file: Arc<dyn File>) -> Result<u64, Errno> {
        let mut fds = self.fds.lock();
        if let Some(i) = fds.iter().position(|f| f.is_none()) {
            fds[i] = Some(file);
            return Ok(i as u64);
        }
        if fds.len() >= MAX_FDS {
            return Err(Errno::EMFILE);
        }
        fds.push(Some(file));
        Ok(fds.len() as u64 - 1)
    }

    pub fn close(&self, fd: u64) -> Result<(), Errno> {
        let file = {
            let mut fds = self.fds.lock();
            match fds.get_mut(fd as usize) {
                Some(slot @ Some(_)) => slot.take(),
                _ => return Err(Errno::EBADF),
            }
        };
        drop(file); // dropping the last Arc closes the file (after unlocking)
        Ok(())
    }

    pub fn status(&self) -> Option<ExitStatus> {
        without_interrupts(|| *self.status.lock())
    }

    /// Make `path` absolute, relative to this process's current directory.
    pub fn resolve(&self, path: &str) -> Result<String, Errno> {
        let cwd = self.cwd.lock().clone();
        fs::normalize(&cwd, path)
    }
}

/// Look up a live (or not yet reaped) process.
pub fn get(pid: u64) -> Option<Arc<Process>> {
    PROCESSES.lock().get(&pid).cloned()
}

/// End the current process (exit syscall, or killed by an exception).
pub fn exit_current(status: ExitStatus) -> ! {
    if let Some(process) = thread::current_process() {
        // Close all files (a pipe's reader sees EOF once its writers are gone).
        let files = core::mem::take(&mut *process.fds.lock());
        drop(files);
        match status {
            ExitStatus::Exited(code) => {
                crate::kprintln!("[process] pid {} ({}) exited with code {}", process.pid, process.name, code)
            }
            ExitStatus::Killed(why) => {
                crate::kprintln!("[process] pid {} ({}) killed: {}", process.pid, process.name, why)
            }
        }
        without_interrupts(|| *process.status.lock() = Some(status));
        process.exited.wake_all(); // the parent may be blocked in wait()
    }
    thread::exit_current()
}

/// Block (the calling thread) until `process` has ended, then remove it from
/// the process table. Used by the kernel and by the `wait` syscall.
pub fn wait(process: &Process) -> ExitStatus {
    let status = process.exited.wait_until(|| *process.status.lock());
    PROCESSES.lock().remove(&process.pid);
    status
}

/// The `wait(pid)` syscall: only a process's parent may wait for it.
pub fn wait_child(parent: u64, pid: u64) -> Result<ExitStatus, Errno> {
    let child = get(pid).filter(|c| c.parent == parent).ok_or(Errno::ECHILD)?;
    Ok(wait(&child))
}