//! Processes = a thread + its own address space + a file descriptor table.
//!
//! Starting one (`spawn`):
//!   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 = the console (tty.rs)
//!   5. spawn a thread with that CR3. The thread starts in the kernel
//!      (`user_thread_entry`) and "returns" into ring 3 with `iretq`.
//!
//! From then on the process only comes back into the kernel through
//! interrupts (timer -> maybe another thread runs), exceptions (page fault in
//! user mode -> the process is killed, see interrupts.rs) and syscalls.
//!
//! Not done yet (exercises): freeing a dead process's frames and page tables,
//! fork/exec/wait syscalls, passing arguments, more than one thread per process.

use crate::elf::{self, ElfError};
use crate::errno::Errno;
use crate::fs::file::File;
use crate::memory::{self, USER_END};
use crate::{serial_println, thread, tty::Console};
use alloc::{sync::Arc, vec, vec::Vec};
use core::sync::atomic::{AtomicU64, Ordering};
use spin::Mutex;
use x86_64::VirtAddr;
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;

#[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),
}

pub struct Process {
    pub pid: u64,
    pub name: &'static str,
    entry: u64,
    fds: Mutex<Vec<Option<Arc<dyn File>>>>,
    status: Mutex<Option<ExitStatus>>,
}

static NEXT_PID: AtomicU64 = AtomicU64::new(1);

/// Load ELF program `elf` into a new address space and start it.
pub fn spawn(name: &'static str, elf_image: &[u8]) -> 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 console: Arc<dyn File> = Arc::new(Console);
    let process = Arc::new(Process {
        pid: NEXT_PID.fetch_add(1, Ordering::Relaxed),
        name,
        entry,
        fds: Mutex::new(vec![Some(console.clone()), Some(console.clone()), Some(console)]),
        status: Mutex::new(None),
    });
    serial_println!("[process] started pid {} ({}), entry {:#x}", process.pid, name, entry);
    thread::spawn_with(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");
    // Functions expect rsp % 16 == 8 on entry (as if a return address had
    // just been pushed), so start 8 bytes below the 16-aligned top.
    unsafe { enter_user_mode(process.entry, USER_STACK_TOP - 8) }
}

/// 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).
unsafe fn enter_user_mode(rip: u64, rsp: 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 esi, esi", "xor edi, edi", "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,
            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 mut fds = self.fds.lock();
        match fds.get_mut(fd as usize) {
            Some(slot @ Some(_)) => {
                *slot = None; // dropping the Arc closes the file
                Ok(())
            }
            _ => Err(Errno::EBADF),
        }
    }

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

/// End the current process (exit syscall, or killed by an exception).
pub fn exit_current(status: ExitStatus) -> ! {
    if let Some(process) = thread::current_process() {
        process.fds.lock().clear(); // close all files
        *process.status.lock() = Some(status);
        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
            ),
        }
    }
    thread::exit_current()
}

/// Block (the calling kernel thread) until `process` has ended.
pub fn wait(process: &Process) -> ExitStatus {
    loop {
        if let Some(status) = process.status() {
            return status;
        }
        x86_64::instructions::hlt(); // the timer lets the process run meanwhile
    }
}