use core::{arch::riscv64::wfi, array, cell::LazyCell, time::Duration};

use alloc::string::String;
use log::info;

use crate::{
    process::{create_process, ExecutionContext, Process, ProcessState},
    time,
};

pub const PROCESS_COUNT: usize = 16;

pub static mut ACTIVE_PID: usize = 0;
pub static mut PROCESS_TABLE: LazyCell<[Process; PROCESS_COUNT]> = LazyCell::new(|| {
    array::from_fn(|_| Process {
        pid: -1,
        name: String::new(),
        state: ProcessState::Dead,
        wake_time: Duration::new(0, 0),
        ctx: ExecutionContext {
            ra: core::ptr::null(),
            sp: core::ptr::null(),
            gp: 0,
            tp: 0,
            a: [0; _],
            t: [0; _],
            s: [0; _],
            mepc: core::ptr::null(),
            mstatus: 0,
        },
        stack: [0; _],
        entry: None,
    })
});

pub fn idle() {
    loop {
        // write_string_temp("idle");
        // info!("idle");
        unsafe {
            wfi();
        }
    }
}

pub fn scheduler_init() {
    info!("scheduler init");
    for pid in 0..PROCESS_COUNT {
        unsafe {
            PROCESS_TABLE[pid].state = ProcessState::Dead;
        }
    }

    create_process(&idle, "idle");
    unsafe {
        PROCESS_TABLE[0].state = ProcessState::Active;
    }
}

pub fn scheduler_without_ret(interrupt_state: &mut *mut ExecutionContext) {
    // info!("scheduler");
    unsafe {
        let prev_pid = ACTIVE_PID;
        PROCESS_TABLE[prev_pid].ctx = **interrupt_state;

        if PROCESS_TABLE[prev_pid].state == ProcessState::Active {
            PROCESS_TABLE[prev_pid].state = ProcessState::Activable;
        }

        loop {
            if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Asleep
                && time::elapsed_time_since_startup() > PROCESS_TABLE[ACTIVE_PID].wake_time
            {
                PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Activable;
            }
            ACTIVE_PID = (ACTIVE_PID + 1) % PROCESS_COUNT;
            if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Activable {
                break;
            }
        }
        PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Active;

        *interrupt_state = &raw mut PROCESS_TABLE[ACTIVE_PID].ctx
    }
}