#define MODULE "gdt"

#include "gdt/GDT.h"
#include "log/Log.h"
#include "memory/MemoryManager.h"
#include "std/assert.h"
#include "std/string.h"
#include <stdint.h>

struct GDTR
{
    uint16_t size;
    uint64_t offset;
} __attribute__((packed));

struct GDTEntry
{
    uint16_t limit0;
    uint16_t base0;
    uint8_t base1;
    uint8_t access;
    uint8_t limit1_flags;
    uint8_t base2;
} __attribute__((packed));

struct HighGDTEntry
{
    uint32_t base_high;
    uint32_t reserved;
} __attribute__((packed));

struct TSS
{
    uint32_t reserved0;
    uint64_t rsp[3];
    uint64_t reserved1;
    uint64_t ist[7];
    uint64_t reserved2;
    uint16_t reserved3;
    uint16_t iomap_base;
} __attribute__((packed));

struct InternalGDT
{
    GDTEntry null;
    GDTEntry kernel_code;
    GDTEntry kernel_data;
    GDTEntry user_code;
    GDTEntry user_data;
    GDTEntry tss;
    HighGDTEntry tss2;
} __attribute__((packed)) __attribute((aligned(PAGE_SIZE)));

__attribute__((aligned(PAGE_SIZE))) static InternalGDT internal_gdt = {{0x0000, 0x0000, 0x00, 0x00, 0x00, 0x00},
                                                                       {0xffff, 0x0000, 0x00, 0x9a, 0xaf, 0x00},
                                                                       {0xffff, 0x0000, 0x00, 0x92, 0xcf, 0x00},
                                                                       {0xffff, 0x0000, 0x00, 0xfa, 0xaf, 0x00},
                                                                       {0xffff, 0x0000, 0x00, 0xf2, 0xcf, 0x00},
                                                                       {0x0000, 0x0000, 0x00, 0xe9, 0x0f, 0x00},
                                                                       {0x00000000, 0x00000000}};

static TSS main_tss;

extern "C" void load_gdt(GDTR* gdtr);
extern "C" void load_tr(int segment);

static void set_base(GDTEntry* entry, uint32_t base)
{
    entry->base0 = (base & 0xFFFF);
    entry->base1 = (base >> 16) & 0xFF;
    entry->base2 = (uint8_t)((base >> 24) & 0xFF);
}

static void set_limit(GDTEntry* entry, uint32_t limit)
{
    ASSERT(limit <= 0xFFFFF);
    entry->limit0 = limit & 0xFFFF;
    entry->limit1_flags = (entry->limit1_flags & 0xF0) | ((limit >> 16) & 0xF);
}

static void set_tss_base(GDTEntry* tss1, HighGDTEntry* tss2, uint64_t addr)
{
    set_base(tss1, addr & 0xffffffff);
    tss2->base_high = (uint32_t)(addr >> 32);
}

static void setup_tss()
{
    memset(&main_tss, 0, sizeof(TSS));
    main_tss.rsp[0] =
        (uint64_t)MemoryManager::get_pages(4) + (PAGE_SIZE * 4) - 8; // allocate 16KB for the syscall stack
    main_tss.iomap_base = sizeof(TSS);
    set_tss_base(&internal_gdt.tss, &internal_gdt.tss2, (uint64_t)&main_tss);
    set_limit(&internal_gdt.tss, sizeof(TSS) - 1);
}

static void load_tss()
{
    kdbgln("Loading TR (GDT entry 0x2b)");
    load_tr(0x2b);
}

void GDT::load()
{
    static GDTR gdtr;
    gdtr.offset = (uint64_t)&internal_gdt;
    gdtr.size = sizeof(InternalGDT);
    setup_tss();
    kdbgln("Loading GDT at offset %lx, size %d", gdtr.offset, gdtr.size);
    load_gdt(&gdtr);
    load_tss();
}