//!
//! Scheduler and idle loop utilities.
//!
//! This module exposes the global process table, the scheduler initialization
//! and a simple round-robin scheduler used by the kernel.
use core::{arch::riscv64::wfi, cell::LazyCell, ops::Bound};

use alloc::{boxed::Box, collections::BTreeMap};
use log::info;

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

#[derive(Debug)]
pub struct Scheduler {
    pub next_pid: u64,
    pub active_pid: u64,
    pub process_table: BTreeMap<u64, Box<Process>>,
}

pub static SCHEDULER: Mutex<LazyCell<Scheduler>> = Mutex::new(LazyCell::new(|| Scheduler {
    next_pid: 0,
    active_pid: 0,
    process_table: BTreeMap::new(),
}));

/// Idle loop executed when there is no runnable process.
///
/// Uses the `wfi` instruction to yield the CPU while waiting for interrupts.
pub fn idle() {
    loop {
        // write_string_temp("idle");
        // info!("idle");
        unsafe {
            wfi();
        }
    }
}

impl Scheduler {
    /// Initialize the scheduler and create the idle process.
    ///
    /// Marks all process slots as `Dead` then creates the idle process and sets
    /// it as the active process.
    pub fn init(&mut self) {
        info!("scheduler init");
        self.create_process(Box::new(idle), "idle");
        self.process_table.get_mut(&0).unwrap().state = ProcessState::Active;
    }

    /// Round-robin scheduler used to select the next runnable process.
    ///
    /// Saves the provided interrupt context into the previous process slot and
    /// updates `ACTIVE_PID` to point to the chosen process. This function does
    /// not return but instead updates `interrupt_state` to the context of the
    /// next process to run.
    pub fn schedule(&mut self, interrupt_state: &mut *mut ExecutionContext) {
        // info!("scheduler");
        unsafe {
            let prev_pid = self.active_pid;
            if let Some(previous_process) = self.process_table.get_mut(&prev_pid) {
                previous_process.ctx = **interrupt_state;

                if previous_process.state == ProcessState::Active {
                    previous_process.state = ProcessState::Activable;
                }
            }

            let mut current_process_iter = self
                .process_table
                .range_mut((Bound::Excluded(prev_pid), Bound::Unbounded));

            self.active_pid = loop {
                if let Some((pid, current_process)) = current_process_iter.next() {
                    if current_process.state == ProcessState::Asleep
                        && time::elapsed_time_since_startup() > current_process.wake_time
                    {
                        current_process.state = ProcessState::Activable;
                    }
                    if current_process.state == ProcessState::Activable {
                        current_process.state = ProcessState::Active;
                        *interrupt_state = &raw mut current_process.ctx;
                        break *pid;
                    };
                } else {
                    current_process_iter = self
                        .process_table
                        .range_mut((Bound::Unbounded, Bound::Included(prev_pid)))
                }
            };
        }
    }
    pub fn get_current_process(&mut self) -> &mut Process {
        self.process_table.get_mut(&self.active_pid).unwrap()
    }
}