#include "tasking.hpp"

#include <array>

struct Task
{
    enum class State {
        Invalid,
        Staging,
        Staged,
        Running
    };

    using enum State;

    std::uint32_t esp;
    std::uint32_t ebp;
    State state = State::Invalid;
};

static std::array<Task, 4> tasks;
static int current = -1;

void schedule(const Registers& regs)
{
    if (current < 0)
        return;

    asm volatile(R"(
        mov %%esp, %0
        mov %%ebp, %1
    )" : "=m" (tasks[current].esp), "=m" (tasks[current].ebp));

    do {
        if (++current >= tasks.size())
            current = 0;
    } while (tasks[current].state == Task::Invalid || tasks[current].state == Task::Staging);

    asm volatile(R"(
        mov %0, %%esp
        mov %1, %%ebp
    )" :: "m" (tasks[current].esp), "m" (tasks[current].ebp));

    if (tasks[current].state == Task::Staged) {
        tasks[current].state = Task::Running;
        asm volatile(R"(
            pop %eax
            popa
            add $0x4, %esp
            iret
        )");
    }
}

void tasking_initialize()
{
    tasks[0].state = Task::Running;
    current = 0;
    asm volatile("int $0x20");
}

bool tasking_spawn(void (*entry)(), unsigned ssize)
{
    int i = -1;
    for (i = 0; i < tasks.size(); ++i) {
        if (tasks[i].state == Task::Invalid)
            break;
    }

    if (i < 0)
        return false;

    tasks[i].state = Task::Staging;

    auto stack = reinterpret_cast<std::uint32_t>(new std::uint8_t[ssize]);
    const auto stackend = stack + ssize;
    const auto regbase = stackend - sizeof(Registers);
    auto r = reinterpret_cast<Registers *>(regbase);
    r->ebp = stackend;
    r->esp = stackend;
    r->eip = reinterpret_cast<std::uint32_t>(entry);
    r->cs = 0x8;
    asm volatile("pushfl; pop %%eax" : "=a"(r->eflags));

    tasks[i] = Task {
        .esp = regbase,
        .ebp = stackend,
        .state = Task::Staged
    };
    return true;
}