/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   Global Descriptor Table and Task State Segment setup
 * @license GPL-3.0-only
 */

#include <arch/gdt.h>
#include <stdint.h>
#include <io/serial/serial.h>
#include <kernel.h>

// Descriptors are 8-byte wide (64bits)
// So the selectors will be (in bytes): 0x0, 0x8, 0x10, 0x18, etc..
uint64_t gdt_entries[NUM_GDT_ENTRIES];
struct GDTR gdtr;

struct tss tss = {0};

/*
 * gdt_load - Loads Global Descriptor Table
 */
static void gdt_load()
{
    asm("lgdt %0" : : "m"(gdtr));
}

/*
 * gdt_flush - Flushes the Global Descriptor Table
 *
 * This function loads new Segment Selectors to make
 * the GDT changes take effect
 */
static void gdt_flush()
{
    // Here, 0x8 is the kernel code selector
    // and 0x10 is the kernel data selector
    asm volatile (
        "mov $0x10, %%ax \n" // Reload segments with kernel data selector
        "mov %%ax, %%ds \n"
        "mov %%ax, %%es \n"
        "mov %%ax, %%fs \n"
        "mov %%ax, %%gs \n"
        "mov %%ax, %%ss \n"

        "pushq $0x8      \n"   // CS reload
        "lea 1f(%%rip), %%rax \n"
        "push %%rax     \n"
        "lretq          \n"
        "1:             \n" // Execution continues here after CS reload
        :
        :
        : "rax", "memory"
    );
}

/*
 * get_set_entry - Sets a GDT entry
 * @num:    Number of the entry (index in GDT)
 * @flags:  Flags (Granularity, Size, Long mode)
 * @access: Access byte (contains Descriptor Privilege Level)
 *
 * This function fills a GDT entry with the specified @flags
 * and @access byte. The base and limit fields are left to zero
 * because we don't use segmentation for memory management.
 */

static void gdt_set_entry(int num, uint8_t flags, uint8_t access)
{
    uint64_t gdt_entry = 0;

    gdt_entry |= (access << 8);
    gdt_entry |= (flags << 20);
    
    // Rest (base, limit) is always zero
    gdt_entries[num] = gdt_entry << 32;
}

/*
 * gdt_set_tss - Setup the TSS entry in the GDT
 * @num: Number of the entry (index in GDT)
 *
 * This function sets up a Task State Segment entry
 * in the Global Descriptor Table.
 *
 * The entry is 128-bit long, so it actually takes
 * two 64-bit GDT entries.
 */
static void gdt_set_tss(int num)
{
    uint64_t tss_base = (uint64_t)&tss;
    uint64_t tss_limit = sizeof(struct tss) - 1;

    tss.iopb = sizeof(struct tss);

    uint64_t tss_low = 0;
    tss_low |= (tss_limit & 0xFFFFULL);
    tss_low |= (tss_base & 0xFFFFFFULL) << 16;
    tss_low |= 0x89ULL << 40;
    tss_low |= ((tss_limit >> 16) & 0xFULL) << 48;
    tss_low |= ((tss_base >> 24) & 0xFFULL) << 56;

    uint64_t tss_high = (tss_base >> 32) & 0xFFFFFFFFULL;

    gdt_entries[num] = tss_low;
    gdt_entries[num + 1] = tss_high;
}

/*
 * gdt_init - Global Descriptor Table initialization 
 *
 * This function loads a new GDT in the CPU.
 * It contains a null descriptor, kernel code and data
 * segments, and user code and data segments.
 * However, we do not use segmentation to manage memory on
 * 64-bit x86, as it's deprecated. Instead, we use paging. 
 */
void gdt_init()
{
    gdt_set_entry(0, 0, 0); // Null descriptor  (0x0)
    gdt_set_entry(1, 0xA, 0x9B); // Kernel code (0x8)
    gdt_set_entry(2, 0xC, 0x93); // Kernel data (0x10)
    gdt_set_entry(3, 0xA, 0xFB); // User code   (0x18)
    gdt_set_entry(4, 0xC, 0xF3); // User data   (0x20)
    gdt_set_tss(5);                            // TSS         (0x28)

    // The -1 subtraction is some wizardry explained in the OSDev wiki -> GDT
    gdtr.limit = NUM_GDT_ENTRIES * sizeof(uint64_t) - 1;
    gdtr.address = (uint64_t)gdt_entries;

    // Load the GDT we created, flush the old one
    gdt_load();
    gdt_flush();

    // Load task register with new TSS
    asm volatile("ltr %%ax" : : "a"(TSS_SEGMENT) : "memory");

    DEBUG("GDT initialized");
}