//! The scheduler that coordinates the execution of threads

use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering::Relaxed;

use arraydeque::ArrayDeque;

use crate::arch::int::lapic::{IPIDestination, LAPIC};
use crate::arch::{cpu, int};
use crate::interrupts::guard::GUARD;
use crate::interrupts::Vector;
use crate::util::{Align, PerCPU};
use crate::{INIT_STACKS, MAX_CPUS};

use super::Thread;

pub const APPS: usize = 10;

static mut THREAD_STACKS: [[u32; 0x4000]; APPS] = [[0; 0x4000]; APPS];

/// The scheduler plans the threads' execution order and, from this,
/// selects the next thread to be running.
///
/// The scheduler manages the ready queue,
/// that is the list of threads that are ready to execute. The scheduler
/// arranges threads in a FIFO order, that is, when a thread is set ready, it
/// will be appended to the end of the queue, while threads to be executed are
/// taken from the front of the queue.
pub struct Scheduler {
    /// Ready queue
    ready: ArrayDeque<usize, APPS>,
    /// Contains all threads
    ///
    /// Note: the idle threads have the id and index of their corresponding CPU
    threads: [Option<Thread>; MAX_CPUS + APPS],
    /// Next thread id
    next_id: AtomicUsize,
    /// Per-CPU data
    local: PerCPU<Local>,
}

/// Private, CPU-local data
struct Local {
    /// Active thread (per core)
    active: Option<usize>,
    /// Thread to be removed after dispatch
    exited: Option<usize>,
}

impl Scheduler {
    /// Construct the scheduler.
    pub const fn new() -> Self {
        Self {
            ready: ArrayDeque::new(),
            threads: [const { None }; MAX_CPUS + APPS],
            next_id: AtomicUsize::new(MAX_CPUS),
            local: PerCPU::new(
                [const {
                    Align(Local {
                        active: None,
                        exited: None,
                    })
                }; MAX_CPUS],
            ),
        }
    }

    /// Add and ready a new thread to the scheduler.
    pub fn add(&mut self, mut thread: Thread) -> usize {
        let id = self.next_id.fetch_add(1, Relaxed);
        thread.id = id;
        thread.init(unsafe { &mut THREAD_STACKS[id - MAX_CPUS] });

        self.threads[id] = Some(thread);
        self.ready(id);
        id
    }

    fn is_idle(thread: usize) -> bool {
        thread < MAX_CPUS
    }

    /// Include a thread in scheduling decisions.
    ///
    /// This method will register a thread for scheduling. It will be appended
    /// to the ready queue and dispatched once its time has come.
    ///
    /// Note: New threads have to be added first with [Scheduler::add].
    pub fn ready(&mut self, thread: usize) {
        assert!(thread < self.threads.len() && self.threads[thread].is_some());
        self.ready.push_back(thread).unwrap();
    }

    pub fn active(&self) -> Option<usize> {
        self.local.get().active
    }

    pub fn thread(&self, thread: usize) -> Option<&Thread> {
        self.threads[thread].as_ref()
    }
    pub fn thread_mut(&mut self, thread: usize) -> Option<&mut Thread> {
        self.threads[thread].as_mut()
    }

    /// Start the scheduling.
    /// This function does not return.
    pub fn schedule(&mut self) -> ! {
        // Initialize own idle thread
        let idle_id = cpu::id();
        let mut idle = Thread::new(Self::idle_action);
        idle.id = idle_id;
        idle.init(unsafe { &mut INIT_STACKS[idle_id] });

        self.threads[idle_id] = Some(idle);

        let next = self.next();
        self.dispatch(next);

        unreachable!();
    }

    /// Initiates a thread switch.
    /// This function returns then in the context of the next thread.
    ///
    /// If `ready`, the currently running thread is added back to the ready list.
    pub fn resume(&mut self, ready: bool) {
        if let Some(active) = self.active().and_then(|t| self.thread(t)) {
            if active.exited {
                if let Some(exited) = self.local.get().exited.replace(active.id) {
                    // If there is already a dead thread
                    assert_ne!(exited, active.id, "Cannot exit while active");
                    self.threads[exited] = None;
                }
            } else if !Self::is_idle(active.id) && ready {
                self.ready.push_back(active.id).unwrap();
            }
        }

        let next = self.next();
        self.dispatch(next);

        if let Some(thread) = self.local.get().exited.take() {
            self.threads[thread] = None;
        }
    }

    /// Terminates the currently running thread, directly continue with the next one.
    pub fn exit(&mut self) -> ! {
        self.thread_mut(self.active().unwrap()).unwrap().exited = true;
        self.resume(false);

        unreachable!();
    }

    /// Terminates `t`, which might be running or waiting.
    pub fn kill(&mut self, thread: usize) -> bool {
        if Self::is_idle(thread) || self.threads[thread].is_none() {
            return false;
        }
        println!(dbg: "kill {thread}");

        if self.local.get().active == Some(thread) {
            self.exit();
        } else {
            if let Some(thread) = &mut self.threads[thread] {
                // will be removed from ready list automatically
                thread.exited = true;
            } else {
                return true;
            }

            // interrupt the running thread
            for cpu in cpu::iter() {
                if unsafe { self.local.get_raw(cpu).active == Some(thread) } {
                    LAPIC.ipi_send(IPIDestination::Physical(cpu as _), Vector::Assassin);
                    break;
                }
            }
        }
        true
    }

    /// Updates the life pointer to next and issues a thread change from
    /// the old to the new life pointer.
    fn dispatch(&mut self, next: usize) {
        let previous = if let Some(previous) = self.local.get().active.replace(next) {
            // Usually we cannot have two mutable references, but I'm sure its ok here
            // -> Lifetime hack: cast it into a pointer and dereference it again
            Some(unsafe { &mut *(self.thread_mut(previous).unwrap() as *mut Thread) })
        } else {
            None
        };

        let next = self.thread(next).unwrap();

        if let Some(last) = previous {
            last.resume(next);
        } else {
            next.go();
        }
    }

    /// Helper to retrieve next Thread.
    fn next(&mut self) -> usize {
        while let Some(thread_id) = self.ready.pop_front() {
            if let Some(thread) = &self.threads[thread_id] {
                if thread.exited {
                    // remove exited threads
                    self.threads[thread_id] = None;
                } else {
                    return thread.id;
                }
            }
        }
        // idle thread
        cpu::id()
    }

    /// Action of the idle thread.
    pub extern "C" fn idle_action() {
        loop {
            int::enable(false);
            if unsafe { GUARD.read().scheduler.ready.is_empty() } {
                let no_timer = cpu::id() != 0 && unsafe { GUARD.read().bell_ringer.is_empty() };
                if no_timer {
                    LAPIC.timer_enable(false);
                }

                // enable int and wait
                int::idle();

                if no_timer {
                    LAPIC.timer_enable(true);
                }
            } else {
                int::enable(true);
                let mut guarded = GUARD.enter();
                guarded.scheduler.resume(false);
            }
        }
    }
}