#pragma once

#include "memory/pointer.h"
#include "util/optional.h"

struct [[gnu::packed]] gdt_register_t {
    uint16_t limit;
    uint64_t base;
};
static_assert(sizeof(gdt_register_t) == 10);

struct Gdt {
    using selector_t = uint16_t;

    enum class segment_type_t : uint8_t {
        code,
        data,
        task_state,
    };

    struct flags_t {
        segment_type_t type;
        uint8_t privilege_level;
        bool long_mode;
        bool default_operand_32bit;
        bool granularity_4KiB;
        bool readable_writeable;
    };

    struct segment_t {
        flags_t flags;
        uint64_t base;
        uint32_t limit;
    };

    static Gdt fromCurrent();
    static optional<Gdt> allocNew();

    void activate();
    optional<selector_t> add(segment_t segment);
    optional<segment_t> get(selector_t selector);

    paddr_t base;
    uint16_t limit;
};

struct [[gnu::packed]] tss_t {
    uint32_t reserved0;
    uint64_t rsp[3];
    uint64_t reserved1;
    uint64_t ist[7];
    uint64_t reserved2;
    uint16_t reserved3;
    uint16_t iomap_base;
};
static_assert(sizeof(tss_t) == 104);

void reloadSegments(Gdt::selector_t code_selector, Gdt::selector_t data_selector);

[[noreturn]] inline void longjump(
    uint64_t code, uint64_t stack, Gdt::selector_t code_selector, Gdt::selector_t stack_selector, uint64_t arg
) {
    uint64_t code_sel = code_selector;
    uint16_t stack_sel = stack_selector;
    asm volatile(
        "mov %[stack_sel], %%ss\n"
        "mov %[stack], %%rsp\n"
        "pushq %[code_sel]\n"
        "pushq %[code]\n"
        "mov %[arg], %%rdi\n"
        "lretq\n"
        :
        : [stack_sel] "r"(stack_sel), [stack] "r"(stack), [code_sel] "r"(code_sel), [code] "r"(code), [arg] "r"(arg)
        : "memory", "rdi"
    );
    __builtin_unreachable();
}