#include "multitasking.hpp"

#include <new>

extern "C" void taskEntryThunk();
extern "C" void taskEntry(Task *task);
void taskEntryCxx(Task *task);

void
Task::init(GlobalDescriptorTable &gdt)
{
    cpuState = new(stack + sizeof(stack) - sizeof(CPUState)) CPUState;

    cpuState->eax() = reinterpret_cast<std::uint32_t>(this);
    cpuState->ebx() = 0;
    cpuState->ecx() = 0;
    cpuState->edx() = 0;

    cpuState->esi() = 0;
    cpuState->edi() = 0;
    cpuState->ebp() = 0;

    cpuState->gs() = gdt.getDS();
    cpuState->fs() = gdt.getDS();
    cpuState->es() = gdt.getDS();
    cpuState->ds() = gdt.getDS();

    // cpuState->esp() = ;
    cpuState->eip() = reinterpret_cast<std::uint32_t>(&taskEntryThunk);
    cpuState->cs() = gdt.getCS();
    // cpuState->ss() = ;
    cpuState->eflags() = 0x202;
}

void
taskEntry(Task *task)
{
    taskEntryCxx(task);
}

void
taskEntryCxx(Task *task)
{
    task->entry();
}

void
Task::entry()
{
    (*task)();
    done = true;
    asm ("int $0x20");
}

TaskManager::TaskManager(GlobalDescriptorTable &gdt)
    : InterruptHandler(0x20), gdt(gdt), numTasks(0), currentTask(-1)
{
}

bool
TaskManager::addTask(std::unique_ptr<Task> task)
{
    if (numTasks == 256) {
        return false;
    }

    task->init(gdt);
    tasks[numTasks++] = std::move(task);
    return true;
}

std::uint32_t
TaskManager::handleInterrupt(std::uint32_t esp)
{
    auto cpuState = new(reinterpret_cast<void *>(esp)) CPUState;
    return reinterpret_cast<std::uint32_t>(schedule(cpuState));
}

CPUState *
TaskManager::schedule(CPUState *cpuState)
{
    if (numTasks == 0) {
        return cpuState;
    }

    if (currentTask >= 0) {
        tasks[currentTask]->cpuState = cpuState;
        if (tasks[currentTask]->done) {
            tasks[currentTask] = std::move(tasks[numTasks - 1]);
            --numTasks;
            currentTask %= numTasks;
            // XXX: this could be the last task, need to handle this.
            return tasks[currentTask]->cpuState;
        }
    }

    if (++currentTask >= numTasks) {
        currentTask %= numTasks;
    }
    return tasks[currentTask]->cpuState;
}