#include <stdint.h>

#include "bits.h"
#include "gdt.h"

#define GDT_ENTRIES 5

#define GDTENTRY_G BIT(7)
#define GDTENTRY_DB BIT(6)
#define GDTENTRY_P BIT(7)
#define GDTENTRY_DPL(x) (x << 5)
#define GDTENTRY_S BIT(4)

#define GDTENTRY_TYPE_CODE_DATA(x) (x << 3)
#define GDTENTRY_TYPE_READ BIT(1)

#define GDTENTRY_RING0_CODE (GDTENTRY_P | GDTENTRY_DPL(0) | GDTENTRY_S | GDTENTRY_TYPE_CODE_DATA(1) | GDTENTRY_TYPE_READ)
#define GDTENTRY_RING3_CODE (GDTENTRY_P | GDTENTRY_DPL(3) | GDTENTRY_S | GDTENTRY_TYPE_CODE_DATA(1) | GDTENTRY_TYPE_READ)

#define GDTENTRY_RING0_DATA (GDTENTRY_P | GDTENTRY_DPL(0) | GDTENTRY_S | GDTENTRY_TYPE_CODE_DATA(0) | GDTENTRY_TYPE_READ)
#define GDTENTRY_RING3_DATA (GDTENTRY_P | GDTENTRY_DPL(3) | GDTENTRY_S | GDTENTRY_TYPE_CODE_DATA(0) | GDTENTRY_TYPE_READ)

extern void reload_segments();

typedef struct {
    uint16_t limit_low;
    uint16_t base_low;
    uint8_t base_middle;
    uint8_t access;
    uint8_t flags_limit_high;
    uint8_t base_high;
} gdt_entry_t;

typedef struct {
    uint16_t size;
    uint32_t address;
} __attribute__((packed)) gdt_desc_t;

static gdt_entry_t gdt[GDT_ENTRIES] = { 0 };

static gdt_desc_t gdt_desc = {
    .size = sizeof(gdt) - 1,
    .address = (uintptr_t)&gdt[0]
};

static void gdt_populate() {
    // Kernel Code
    gdt[1].base_low = 0;
    gdt[1].base_middle = 0;
    gdt[1].base_high = 0;
    gdt[1].limit_low = 0xFFFF;
    gdt[1].flags_limit_high = 0x0F | 
                              GDTENTRY_G |
                              GDTENTRY_DB; // 4KB granularity, 32-bit segment
    gdt[1].access = GDTENTRY_RING0_CODE;

    // Kernel Data
    gdt[2].base_low = 0;
    gdt[2].base_middle = 0;
    gdt[2].base_high = 0;
    gdt[2].limit_low = 0xFFFF;
    gdt[2].flags_limit_high = 0x0F | 
                              GDTENTRY_G |
                              GDTENTRY_DB; // 4KB granularity, 32-bit segment
    gdt[2].access = GDTENTRY_RING0_DATA;

    // User Code
    gdt[3].base_low = 0;
    gdt[3].base_middle = 0;
    gdt[3].base_high = 0;
    gdt[3].limit_low = 0xFFFF;
    gdt[3].flags_limit_high = 0x0F | 
                              GDTENTRY_G |
                              GDTENTRY_DB; // 4KB granularity, 32-bit segment
    gdt[3].access = GDTENTRY_RING3_CODE;

    // User Data
    gdt[4].base_low = 0;
    gdt[4].base_middle = 0;
    gdt[4].base_high = 0;
    gdt[4].limit_low = 0xFFFF;
    gdt[4].flags_limit_high = 0x0F | 
                              GDTENTRY_G |
                              GDTENTRY_DB; // 4KB granularity, 32-bit segment
    gdt[4].access = GDTENTRY_RING3_DATA;
}

static void gdt_load(const gdt_desc_t *desc) {
    __asm__ volatile ("lgdt %0" :: "m"(*desc));
}

void gdt_init() {
    gdt_populate();
    gdt_load(&gdt_desc);
    reload_segments();
}