mod irq_handlers;
pub mod mutex;

use super::{critical, queue::Queue};
use core::sync::atomic::{compiler_fence, Ordering};
use gba::prelude::{Halt, IME};
use gba::{info, warning};
use heapless::Vec;
use mutex::BlockingMutex;

use crate::{critical_enter_from_thumb, critical_exit_from_thumb};
pub use irq_handlers::*;

#[repr(u32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockingOn {
    Nothing = 0,
    Timer0,
    Timer1,
    Timer2,
    Timer3,
    SerialReadable,
    SerialWritable,
    Mutex,
    VBlank,
    VCount,
}

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Context {
    pub psr: u32,
    pub pc: *const (),
    pub r0: u32,
    pub r1: u32,
    pub r2: u32,
    pub r3: u32,
    pub r4: u32,
    pub r5: u32,
    pub r6: u32,
    pub r7: u32,
    pub r8: u32,
    pub r9: u32,
    pub r10: u32,
    pub r11: u32,
    pub r12: u32,
    pub sp: u32,
    pub lr: u32,
}

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Task {
    /* Task context (registers, psr) */
    pub context: Context,
    /* Condition on which we are blocking. */
    pub block_on: BlockingOn,
    /* Entrypoint. */
    pub entry_addr: extern "C" fn(),
    /* 8-byte-aligned stack, 8 * 512 = 4 kilobytes. */
    pub stack: [u64; 512],
}

impl Context {
    pub const fn from_entrypoint(ep: *const ()) -> Self {
        Context {
            psr: 0x3F, // system mode, thumb, irq/fiq enabled
            pc: ep,
            r0: 0,
            r1: 0x01010101,
            r2: 0x02020202,
            r3: 0x03030303,
            r4: 0x04040404,
            r5: 0x05050505,
            r6: 0x06060606,
            r7: 0x07070707,
            r8: 0x08080808,
            r9: 0x09090909,
            r10: 0x10101010,
            r11: 0x11111111,
            r12: 0x12121212,
            sp: 0x13131313,
            lr: 0x14141414,
        }
    }
}

impl Task {
    pub const fn new(entrypoint: extern "C" fn()) -> Self {
        Task {
            context: Context::from_entrypoint(entrypoint as *const ()),
            block_on: BlockingOn::Nothing,
            entry_addr: entrypoint,
            stack: [0; 512],
        }
    }
}

// Q_PENDING currently requires something like .swap_remove(), which Queue doesn't
// implement yet, and it doesn't require the FIFO behavior that Queue provides,
// so we can just use heapless::Vec.
pub static mut Q_PENDING: Vec<&'static mut Task, 9> = Vec::new();
static mut Q_RUNNABLE: Queue<&'static mut Task, 8> = Queue::new();

static mut CTASK_HOLDER: Option<&mut Task> = None;

#[no_mangle]
static mut CURRENT_TASK: *mut Task = 0 as _;

#[no_mangle]
static mut RDY_TASK: *mut Task = 0 as _;

static mut IDLE_TASK: Task = Task::new(idle_entry);

extern "C" {
    /* Start multitasking, i.e. start the first task, never returning to the caller. */
    fn start_multitasking_from_thumb() -> !;

    /* Low level task switch. */
    fn task_switch_from_thumb();
}

/* Enter a critical section, disabling interrupts and returning the
* processor status register to restore at the end of the section. */
#[instruction_set(arm::a32)]
pub fn _critical_enter() -> u32 {
    unsafe {
        let res: u32;
        asm!(
        "mrs {0}, cpsr",
        "mov {1}, {0}",
        "orr {1}, #0xc0",
        "msr cpsr_c, {1}",
        out(reg) res,
        out(reg) _,
        options(nomem, nostack)
        );
        res
    }
}

/* Exit a critical section, restoring the given PSR. */
#[instruction_set(arm::a32)]
pub fn _critical_exit(psr: u32) {
    unsafe {
        asm!(
        "msr cpsr_c, {}",
        in(reg) psr,
        options(nomem, nostack)
        );
    }
}

#[macro_export]
macro_rules! critical {
    ($block:block) => {
        use core::sync::atomic::{compiler_fence, Ordering};
        use crate::task::{_critical_enter, _critical_exit};
        let psr = _critical_enter();
        compiler_fence(Ordering::SeqCst);
        $block
        compiler_fence(Ordering::SeqCst);
        _critical_exit(psr);
    };
    (unsafe $block:block) => {
        use core::sync::atomic::{compiler_fence, Ordering};
        use crate::task::{_critical_enter, _critical_exit};
        let psr = _critical_enter();
        compiler_fence(Ordering::SeqCst);
        unsafe { $block }
        compiler_fence(Ordering::SeqCst);
        _critical_exit(psr);
    };
}

pub fn register(task: &'static mut Task) {
    critical!(unsafe {
        task.context.sp = task.stack.as_mut_ptr().add(task.stack.len()) as u32;
        Q_RUNNABLE.push(task).unwrap();
    });
}

pub fn block_on(condition: BlockingOn) {
    critical!(unsafe {
        (*CURRENT_TASK).block_on = condition;
        Q_PENDING.push(CTASK_HOLDER.take().unwrap()).unwrap();
        compiler_fence(Ordering::SeqCst);
        schedule::<true>();
    });
}

/* Unblock tasks pending on the given condition. */
pub fn unblock_on(condition: BlockingOn) {
    critical!(unsafe {
        let mut idxs = Vec::<usize, 9>::new();

        for (i, task) in Q_PENDING.iter_mut().enumerate() {
            if task.block_on == condition {
                task.block_on = BlockingOn::Nothing;
                idxs.push(i);
            }
        }

        while let Some(i) = idxs.pop() {
            // This works because indices are in descending order.
            let task = Q_PENDING.swap_remove(i);
            Q_RUNNABLE.push(task).unwrap();
        }
    });
}

/*
* Schedule a new task, swapping out the current one.
* Note that the current task is not put back on a runqueue - the calling function should do this.
* N.B. task_switch sets CURRENT_TASK to RDY_TASK
*/
pub fn schedule<const SWITCH: bool>() {
    critical!(unsafe {
        // Only switch if there is no current task, the task is blocked, or there is another task which can be run.
        if CURRENT_TASK.is_null()
            || (*CURRENT_TASK).block_on != BlockingOn::Nothing
            || !Q_RUNNABLE.is_empty()
        {
            let rdy_task = Q_RUNNABLE.pop().unwrap_or(&mut IDLE_TASK);
            RDY_TASK = rdy_task as *mut _;

            let stack_top = rdy_task.stack.as_mut_ptr().add(rdy_task.stack.len()) as u32;
            let stack_used = stack_top - rdy_task.context.sp;

            info!(
                "RDY_TASK: {:x} stack base: {:x} stack pointer: {:x} stack usage: {}",
                RDY_TASK as u32, stack_top, rdy_task.context.sp, stack_used
            );

            CTASK_HOLDER = Some(rdy_task);

            if SWITCH {
                task_switch_from_thumb(); // switch out the current task for the ready task.
            }
        }
    });
}

pub fn start() -> ! {
    critical!({
        // initialize SP correctly for the idle task
        unsafe {
            IDLE_TASK.context.sp = IDLE_TASK.stack.as_mut_ptr().add(IDLE_TASK.stack.len()) as u32;
        }

        // schedule the first task to run
        schedule::<false>();
        unsafe {
            start_multitasking_from_thumb();
        }
    });
}

extern "C" fn idle_entry() {
    loop {
        // TODO: Halt() here is not right, just use proper interrupt-returns-to-runnable-task
        schedule::<true>();
    }
}
