#include <kernel/tty.h>
#include <debugging.h>
#include <stdint.h>
#include <stddef.h>

#define NUM_ENTRIES 6

struct SegmentDescriptor {
    uint32_t limit;
    uint32_t base;
    uint8_t access_byte;
    uint8_t flags;
};

struct GDT {
    uint16_t size;
    uint64_t* base;
} __attribute__((packed));

struct tss_entry_struct {
	uint32_t prev_tss; // The previous TSS - with hardware task switching these form a kind of backward linked list.
	uint32_t esp0;     // The stack pointer to load when changing to kernel mode.
	uint32_t ss0;      // The stack segment to load when changing to kernel mode.
	// Everything below here is unused.
	uint32_t esp1; // esp and ss 1 and 2 would be used when switching to rings 1 or 2.
	uint32_t ss1;
	uint32_t esp2;
	uint32_t ss2;
	uint32_t cr3;
	uint32_t eip;
	uint32_t eflags;
	uint32_t eax;
	uint32_t ecx;
	uint32_t edx;
	uint32_t ebx;
	uint32_t esp;
	uint32_t ebp;
	uint32_t esi;
	uint32_t edi;
	uint32_t es;
	uint32_t cs;
	uint32_t ss;
	uint32_t ds;
	uint32_t fs;
	uint32_t gs;
	uint32_t ldt;
	uint16_t trap;
	uint16_t iomap_base;
} __attribute__((packed));

struct tss_entry_struct tss;

static uint64_t entries[NUM_ENTRIES];
static size_t current_entry = 0;
struct GDT gdt;

static void encodeGdtEntry(uint8_t *target, struct SegmentDescriptor source)
{
    // Check the limit to make sure that it can be encoded
    if (source.limit > 0xFFFFF) {
        terminal_writestring("WARNING: GDT entry has a limit (");
        print_hex_bytes(&source.limit, 4);
        terminal_writestring(") that is larger than accepted.");
        return;
    }
    
    // Encode the limit
    target[0] = source.limit & 0xFF;
    target[1] = (source.limit >> 8) & 0xFF;
    target[6] = (source.limit >> 16) & 0x0F;
    
    // Encode the base
    target[2] = source.base & 0xFF;
    target[3] = (source.base >> 8) & 0xFF;
    target[4] = (source.base >> 16) & 0xFF;
    target[7] = (source.base >> 24) & 0xFF;
    
    // Encode the access byte
    target[5] = source.access_byte;
    
    // Encode the flags
    target[6] |= (source.flags << 4);

    print_hex_bytes(target, 8);
}

static void add_entry(struct SegmentDescriptor entry)
{
    if (current_entry == NUM_ENTRIES) {
        terminal_writestring("WARNING: More GDT entries than NUM_ENTRIES.\n");
        return;
    }

    uint64_t descriptor = 0;
    encodeGdtEntry((uint8_t*) &descriptor, entry);

    entries[current_entry] = descriptor;

    current_entry++;
}

static void load_gdt() 
{
    gdt.size = NUM_ENTRIES * 8 - 1;
    gdt.base = entries;

    asm ( "lgdt %0" : : "m"(gdt) );
}

void gdt_init()
{
    // https://wiki.osdev.org/Global_Descriptor_Table

    terminal_initialize();
    struct SegmentDescriptor null;
    null.limit = 0;
    null.base = 0;
    null.access_byte = 0;
    null.flags = 0;

    terminal_writestring("Kernel Null: ");

    add_entry(null);

    terminal_putchar('\n');

    struct SegmentDescriptor kernel_code;
    kernel_code.limit = 0x000003FF;
    kernel_code.base =  0x00000000;
    kernel_code.access_byte = 0b10011010;
    kernel_code.flags = 0b1100;

    terminal_writestring("Kernel Code: ");

    add_entry(kernel_code);
    terminal_putchar('\n');

    struct SegmentDescriptor kernel_data;
    kernel_data.limit = 0x000003FF;
    kernel_data.base =  0x00400000;
    kernel_data.access_byte = 0b10010010;
    kernel_data.flags = 0b1100;

    terminal_writestring("Kernel Data: ");
    add_entry(kernel_data);
    terminal_putchar('\n');

    struct SegmentDescriptor user_code;
    user_code.limit = 0x000003FF;
    user_code.base =  0x00800000;
    user_code.access_byte = 0b11111010;
    user_code.flags = 0b1100;

    terminal_writestring("User Code: ");

    add_entry(user_code);
    terminal_putchar('\n');

    struct SegmentDescriptor user_data;
    user_data.limit = 0x000FFFFF;
    user_data.base =  0x00c00000;
    user_data.access_byte = 0b11110010;
    user_data.flags = 0b1100;

    terminal_writestring("User Data: ");

    add_entry(user_data);
    terminal_putchar('\n');

    struct SegmentDescriptor tss_segment;
    tss_segment.limit = sizeof(tss_segment) - 1;
    tss_segment.base =  (uint32_t) &tss;
    tss_segment.access_byte = 0x89;
    tss_segment.flags = 0b0000;

    terminal_writestring("TSS Segment: ");

    add_entry(tss_segment);
    terminal_putchar('\n');

    load_gdt();
}