#ifndef TOS__MULTITASKING_HPP__
#define TOS__MULTITASKING_HPP__

#include <cstdint>

#include <memory>
#include <type_traits>
#include <utility>

#include "hwcomm/interrupts.hpp"
#include "GDT.hpp"

struct CPUState
{
    std::uint32_t data[16];

    std::uint32_t & eax() { return data[0]; }
    std::uint32_t & ebx() { return data[1]; }
    std::uint32_t & ecx() { return data[2]; }
    std::uint32_t & edx() { return data[3]; }

    std::uint32_t & esi() { return data[4]; }
    std::uint32_t & edi() { return data[5]; }
    std::uint32_t & ebp() { return data[6]; }

    std::uint32_t & gs() { return data[7]; }
    std::uint32_t & fs() { return data[8]; }
    std::uint32_t & es() { return data[9]; }
    std::uint32_t & ds() { return data[10]; }

    std::uint32_t & eip()    { return data[11]; }
    std::uint32_t & cs()     { return data[12]; }
    std::uint32_t & eflags() { return data[13]; }
    std::uint32_t & esp()    { return data[14]; }
    std::uint32_t & ss()     { return data[15]; }
};
static_assert(std::is_pod<CPUState>::value, "CPUState must be POD!");

class Task
{
    friend class TaskManager;
    friend void taskEntryCxx(Task *task);

    struct TaskImplBase
    {
        virtual ~TaskImplBase() = default;

        virtual void operator()() = 0;
    };

    template <typename F>
    struct TaskImpl : TaskImplBase
    {
        F f;

        TaskImpl(F &&f) : f(std::forward<decltype(f)>(f))
        {
        }

        virtual void operator()() override
        {
            f();
        }
    };

public:
    template <typename F>
    Task(F &&f)
        : task(new TaskImpl<F>(std::forward<F>(f))),
          cpuState(nullptr), done(false)
    {
    }

private:
    void init(GlobalDescriptorTable &gdt);
    void entry();

private:
    union {
        std::uint64_t forceAlignment;
        std::uint8_t stack[4096]; // 4 KiB
    };
    std::unique_ptr<TaskImplBase> task;
    CPUState *cpuState;
    bool done;
};

class TaskManager : public hwcomm::InterruptHandler
{
    friend void Task::entry();

public:
    TaskManager(GlobalDescriptorTable &gdt);

public:
    bool addTask(std::unique_ptr<Task> task);

private:
    virtual std::uint32_t handleInterrupt(std::uint32_t esp) override;

private:
    CPUState * schedule(CPUState *cpuState);

private:
    GlobalDescriptorTable &gdt;
    std::unique_ptr<Task> tasks[256];
    int numTasks;
    int currentTask;
};

#endif // TOS__MULTITASKING_HPP__