//! Process management module for the operating system.
//!
//! This module provides the structures and functions necessary to create,
//! manage and schedule processes in the kernel. It defines the `Process` and
//! `ExecutionContext` types and helper functions to create processes from
//! in-memory functions or binaries on the filesystem. The module intentionally
//! keeps unsafe usage localized and documented where raw pointers or transmute
//! are required.

use core::time::Duration;

use alloc::{boxed::Box, format, string::String, vec::Vec};
use bffs::{io::Read, path::Path};
use shared::syscall::exit;

use crate::{
    fs::FILE_SYSTEM,
    println,
    scheduler::{scheduler_without_ret, ACTIVE_PID, PROCESS_COUNT, PROCESS_TABLE},
    time::elapsed_time_since_startup,
};

/// Size of the stack allocated to each process (in 64-bit words).
const STACK_SIZE: usize = 4096;

/// MSTATUS bit to enable supervisor mode interrupts.
const MSTATUS_SPIE: u64 = 1 << 5;

/// MSTATUS bit to set previous privilege mode to supervisor.
const MSTATUS_SPP: u64 = 1 << 1;

/// Represents the state of a process in the system.
#[derive(Debug, PartialEq, Eq)]
pub enum ProcessState {
    /// The process is currently executing.
    Active,
    /// The process is ready to execute and waiting to be scheduled.
    Activable,
    /// The process has terminated and its slot can be reused.
    Dead,
    /// The process is sleeping until a specific wake time.
    Asleep,
}

/// Execution context saved during a context switch.
///
/// This structure contains all RISC-V registers that must be
/// preserved during an interrupt or process switch.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ExecutionContext {
    /// Return address register.
    pub ra: *const u64,
    /// Stack pointer register.
    pub sp: *const u64,
    /// Global pointer register.
    pub gp: u64,
    /// Thread pointer register.
    pub tp: u64,
    /// Argument/return value registers (a0-a7).
    pub a: [u64; 8],
    /// Temporary registers (t0-t6).
    pub t: [u64; 7],
    /// Saved registers (s0-s11).
    pub s: [u64; 12],
    /// Machine exception program counter.
    pub mepc: *const u64,
    /// Machine status register.
    pub mstatus: u64,
}

/// Represents a process in the system.
///
/// Each process has its own execution context, stack,
/// and metadata for scheduling.
pub struct Process {
    /// Unique process identifier.
    pub pid: i64,
    /// Descriptive name of the process.
    pub name: String,
    /// Current state of the process.
    pub state: ProcessState,
    /// Optional entry point for the process code.
    pub entry: Option<&'static dyn Fn()>,
    /// Wake time for sleeping processes.
    pub wake_time: Duration,
    /// Saved execution context.
    pub ctx: ExecutionContext,
    /// Process stack.
    pub stack: [u64; STACK_SIZE],
}

impl core::fmt::Debug for Process {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Process")
            .field("pid", &self.pid)
            .field("name", &self.name)
            .field("state", &self.state)
            .field("wake_time", &self.wake_time)
            .field("ctx", &self.ctx)
            .field("stack", &format!("[_; {}]", STACK_SIZE))
            .finish()
    }
}

/// Creates a process from a binary file.
///
/// # Arguments
///
/// * `path` - Path to the executable binary file.
///
/// # Returns
///
/// Returns the PID of the created process, or -1 on failure.
///
/// # Safety
///
/// This function uses `unsafe` to transmute the file content into an
/// executable function. The binary must be in the correct format and
/// conform to the expected ABI.
/// Create a process from an executable binary located on the filesystem.
///
/// Attempts to open `path`, load its contents into memory and create a new
/// kernel process that will execute the loaded binary. Returns the PID of the
/// created process, or -1 on failure.
pub fn create_process_from_file<'a, T: Into<Path<'a>>>(path: T) -> i64 {
    let path = path.into();
    let name = path.as_str();

    // Open and read the binary file
    let mut bin = FILE_SYSTEM.open_file(path).unwrap();
    let mut content: Vec<u8> = Vec::new();
    bin.read_to_end(&mut content).unwrap();

    println!("Loading binary at address: {:x?}", content.as_ptr());

    // Fallback: treat the file as a raw binary blob and execute in-place
    let entry_point =
        unsafe { core::mem::transmute::<*const u8, extern "C" fn()>(Vec::leak(content).as_ptr()) };
    let wrapper = Box::leak(Box::new(move || {
        entry_point();
    }));

    create_process(wrapper, name)
}

/// Creates a new process with the specified code and name.
///
/// # Arguments
///
/// * `code` - Static reference to the function to execute.
/// * `name` - Name of the process (for identification).
///
/// # Returns
///
/// Returns the PID of the created process, or -1 if the process table is full.
///
/// # Safety
///
/// This function manipulates the global process table and initializes
/// the execution context using unsafe operations.
/// Create a new process from a function pointer.
///
/// The provided `code` function will be executed when the process is first
/// scheduled. Returns the new PID, or -1 if the process table is full.
pub fn create_process<T: Into<String>, F: Fn()>(code: &'static F, name: T) -> i64 {
    // Search for a free slot in the process table
    let mut next_pid = 0;
    while next_pid < PROCESS_COUNT && unsafe { PROCESS_TABLE[next_pid].state != ProcessState::Dead }
    {
        next_pid += 1;
    }

    // Check if a slot is available
    if next_pid >= PROCESS_COUNT {
        return -1; // Process table is full
    }

    // SAFETY: Initializing process in the global table.
    // Access is safe because we verified bounds and found a Dead slot.
    unsafe {
        let process = &mut PROCESS_TABLE[next_pid];

        // Configure process metadata
        process.pid = next_pid as i64;
        process.name = name.into();
        process.state = ProcessState::Activable;
        process.entry = Some(code);

        // Configure execution context
        // a0 contains the pointer to the function to execute
        process.ctx.a[0] = process.entry.as_ref().unwrap_unchecked() as *const &dyn Fn() as u64;

        // mepc points to process_launcher which will call the function
        process.ctx.mepc = process_launcher as *const _;

        // Configure mstatus for supervisor mode with interrupts enabled
        process.ctx.mstatus = MSTATUS_SPP | MSTATUS_SPIE;

        // Initialize stack pointer at the top of the stack
        process.ctx.sp = &raw const process.stack[STACK_SIZE - 1];
    }

    next_pid as i64
}

/// Entry point to launch a new process.
///
/// This function is automatically called during the first scheduling
/// of a process. It executes the process code and calls `exit()`
/// if the code doesn't terminate explicitly.
///
/// # Arguments
///
/// * `code` - Pointer to the function to execute.
///
/// # Safety
///
/// This function must be called with a valid pointer to a function.
/// Internal launcher used as the initial program counter for new processes.
///
/// This function is installed into the process `mepc` so that when the new
/// process is scheduled it will run this launcher which calls the user
/// function and ensures the process exits cleanly.
extern "C" fn process_launcher(code: *const &dyn Fn()) {
    // SAFETY: The code pointer was initialized in create_process
    // and points to a valid function.
    unsafe { (*code)() };

    // If user code didn't exit explicitly, call exit() to clean up the process
    exit();
}

/// Terminates the currently active process.
///
/// This function marks the active process as dead and triggers
/// the scheduler to switch to another process.
///
/// # Arguments
///
/// * `interrupt_context` - Interrupt context for state saving.
///
/// # Note
///
/// This function never returns as it transfers control to
/// another process via the scheduler.
/// Terminate the currently active process and switch to the scheduler.
///
/// Marks the active process as dead and transfers control to the scheduler
/// to select the next runnable process. This function does not return.
pub fn exit_process(interrupt_context: &mut *mut ExecutionContext) {
    // SAFETY: ACTIVE_PID is maintained by the scheduler and is always valid.
    unsafe {
        PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Dead;
    }

    // Transfer control to the scheduler (does not return)
    scheduler_without_ret(interrupt_context)
}

/// Puts the active process to sleep for a specified duration.
///
/// The process will be automatically woken up by the scheduler when
/// the wake time is reached.
///
/// # Arguments
///
/// * `duration` - Duration of the sleep.
/// * `interrupt_context` - Interrupt context for state saving.
///
/// # Note
///
/// This function never returns as it transfers control to
/// another process via the scheduler.
/// Put the active process to sleep for `duration` and schedule the next runnable process.
///
/// The wake time is computed from the current uptime; the scheduler will
/// reactivate the process when the wake time is reached.
pub fn sleep(duration: Duration, interrupt_context: &mut *mut ExecutionContext) {
    // SAFETY: ACTIVE_PID is maintained by the scheduler and is always valid.
    unsafe {
        let process = &mut PROCESS_TABLE[ACTIVE_PID];
        process.wake_time = elapsed_time_since_startup() + duration;
        process.state = ProcessState::Asleep;
    }

    // Transfer control to the scheduler (does not return)
    scheduler_without_ret(interrupt_context)
}