use crate::scheduler::{Task, SCHEDULER};

/// Saves task's state onto stack
macro_rules! save_tasks_state {
    () => {
        "
            // Copy task's state
            push [rsp + 32] // SS
            push [rsp + 32] // RSP
            push [rsp + 32] // RFLAGS
            push [rsp + 32] // CS
            push [rsp + 32] // RIP

            // Save task's registers
            push r15
            push r14
            push r13
            push r12
            push r11
            push r10
            push r9
            push r8
            push rdi
            push rsi
            push rdx
            push rcx
            push rbx
            push rax
        "
    };
}

pub(super) use save_tasks_state;

/// Save provided stack pointer into scheduler's queue
pub extern "C" fn save_and_enqueue(sp: u64) {
    SCHEDULER.lock().enqueue_suspended(sp);
}

// Fetch and load next task
pub unsafe extern "C" fn switch_to_next() -> ! {
    // Fetch next task
    let next = SCHEDULER.lock().pop().expect("no task in the task queue");
    match next {
        Task::Suspended(sp) => asm!(
            // Restore task's registers
            "mov rsp, {}",
            "pop rax",
            "pop rbx",
            "pop rcx",
            "pop rdx",
            "pop rsi",
            "pop rdi",
            "pop r8",
            "pop r9",
            "pop r10",
            "pop r11",
            "pop r12",
            "pop r13",
            "pop r14",
            "pop r15",

            // Copy things up the stack
            "add rsp, 80", // Move 8 bytes above exception stack frame
            "push [rsp - 48]", // SS
            "push [rsp - 48]", // RSP
            "push [rsp - 48]", // RFLAGS
            "push [rsp - 48]", // CS
            "push [rsp - 48]", // RIP

            // Signalise end of the interrupt and return
            "push rax",
            "mov al, 32",
            "out 20h, al",
            "pop rax",
            "iretq",

            in(reg) sp,
            options(noreturn),
        ),
        Task::Spawn {
            function,
            stack_start,
        } => asm!(
            "add rsp, 160", // Move above saved registries
            "mov [rsp], {ip}", // Set task's instruction pointer
            "mov [rsp + 24], {sp}", // Set task's stack pointer

            // Signalise the end of the interrupt and return
            "mov al, 32",
            "out 20h, al",
            "iretq",

            ip = in(reg) function,
            sp = in(reg) stack_start,
            options(noreturn),
        ),
    }
}
