#include <gdt.h>
#include <common.h>
#include <string_utils.h>
#include <serial.h>

static gdt_entry g_gdt[7];  // 5 segments + TSS (2 entries)
static gdt_ptr   g_gdt_ptr;
static tss       g_tss;

extern "C" void lgdt_and_reload(void* gdt_ptr);
extern "C" void tss_load(UINT16 selector);

static void gdt_set_entry(UINT32 index, UINT32 base, UINT32 limit,
                           UINT8 access, UINT8 granularity) {
    g_gdt[index].limit_low   = limit & 0xFFFF;
    g_gdt[index].base_low    = base & 0xFFFF;
    g_gdt[index].base_mid    = (base >> 16) & 0xFF;
    g_gdt[index].access      = access;
    g_gdt[index].granularity  = granularity | ((limit >> 16) & 0x0F);
    g_gdt[index].base_high   = (base >> 24) & 0xFF;
}

static void tss_set_descriptor(UINT32 index) {
    UINT64 base = (UINT64)&g_tss;
    UINT32 limit = sizeof(tss) - 1;

    g_gdt[index].limit_low   = limit & 0xFFFF;
    g_gdt[index].base_low    = base & 0xFFFF;
    g_gdt[index].base_mid    = (base >> 16) & 0xFF;
    g_gdt[index].access      = 0x89;
    g_gdt[index].granularity  = 0x00;
    g_gdt[index].base_high   = (base >> 24) & 0xFF;

    g_gdt[index + 1].limit_low   = (limit >> 16) & 0xFFFF;
    g_gdt[index + 1].base_low    = (base >> 32) & 0xFFFF;
    g_gdt[index + 1].base_mid    = (base >> 48) & 0xFF;
    g_gdt[index + 1].access      = 0;
    g_gdt[index + 1].granularity  = 0;
    g_gdt[index + 1].base_high   = 0;
}

void gdt_init(void) {
    serial_write("GDT: initializing\n");

    for (SSINT32 i = 0; i < 7; i++) {
        g_gdt[i] = {0};
    }

    gdt_set_entry(0, 0, 0, 0, 0);
    gdt_set_entry(1, 0, 0xFFFFF, 0x9A, 0xA0);
    gdt_set_entry(2, 0, 0xFFFFF, 0x92, 0xC0);
    gdt_set_entry(3, 0, 0xFFFFF, 0xFA, 0xA0);
    gdt_set_entry(4, 0, 0xFFFFF, 0xF2, 0xC0);
    tss_set_descriptor(5);

    g_tss = {};
    g_tss.iopb_offset = sizeof(tss);

    g_gdt_ptr.limit = sizeof(g_gdt) - 1;
    g_gdt_ptr.base = (UINT64)&g_gdt[0];

    serial_write("GDT: loading GDT + reloading segments\n");
    lgdt_and_reload(&g_gdt_ptr);
    serial_write("GDT: loading TSS\n");
    tss_load(GDT_TSS);

    serial_write("GDT: done\n");
}

void gdt_set_kernel_stack(UINT64 stack) {
    g_tss.rsp0 = stack;
}