#include <kernel.h>
#include <multiboot.h>
#include <acpi.h>
#include <vga.h>
#include <pci.h>
#include <printf.h>
#include <strcmp.h>
#include <gdt.h>
#include <idt.h>
#include <pic8259.h>
#include <ps2.h>
#include <ide.h>
#include <mem.h>
#include <fat12.h>

struct kernel_context ctx = {0};

struct idt_entry g_idt[256] = {0};
struct idtr g_idtr = {0};

extern uint32_t isr_stub_table[32];
extern void ps2_isr();

#define MAGIC_BREAKPOINT __asm__ volatile("xchgw %bx, %bx");

#define USER_DS 0x23  // User data segment selector
#define USER_CS 0x1B  // User code segment selector

extern struct gdt_entry gdt[6];
struct tss_entry tss;

void tss_flush(uint16_t tss_selector) {
    asm volatile (
        "ltr %0"
        : // no output
        : "r" (tss_selector)
    );
}

void test() {
    printf("\n\nHello from user mode!\n");
    for (;;) {}
}

void switch_to_user_mode(void *user_stack, void *user_entry) {
    asm volatile (
        "movl $0x23, %%eax \n"      // User data segment selector
        "mov %%ax, %%ds \n"
        "mov %%ax, %%es \n"
        "mov %%ax, %%fs \n"
        "mov %%ax, %%gs \n"

        "pushl $0x23 \n"            // User stack segment
        "pushl %0 \n"               // User stack pointer
        "pushfl \n"                 // Push EFLAGS
        "popl %%eax \n"
        "orl $0x200, %%eax \n"      // Enable interrupts (set IF)
        "pushl %%eax \n"

        "pushl $0x1B \n"            // User code segment
        "pushl %1 \n"               // User entry point

        "iret \n"                   // Interrupt return (switches to user mode)
        :
        : "r"(user_stack), "r"(user_entry)
        : "eax"
    );
}

void write_tss(int num, uint16_t ss0, uint32_t esp0) {
    uint32_t base = (uint32_t)&tss;
    uint32_t limit = base + sizeof(tss);

    gdt[num].base_low = base & 0xFFFF;
    gdt[num].base_middle = (base >> 16) & 0xFF;
    gdt[num].base_high = (base >> 24) & 0xFF;
    gdt[num].limit_low = limit & 0xFFFF;
    gdt[num].granularity = (limit >> 16) & 0x0F;
    gdt[num].granularity |= 0x40;
    gdt[num].access = 0x89;

    memset(&tss, 0, sizeof(tss));
    tss.ss0 = ss0;
    tss.esp0 = esp0;
    tss.cs = 0x0b;
    tss.ss = tss.ds = tss.es = tss.fs = tss.gs = 0x13;
}

void kmain(struct multiboot_info *info) {
    clear_screen();
    set_color(VGA_COLOR_CYAN);
// Check if the bootloader gave us the upper and lower memory
    if (!(info->flags & 0x1)) goto halt;

    ctx.multi_mmap = (struct multiboot_mmap_entry *)info->memMapAddress;
    ctx.multi_mmap_size = info->memMapLength / sizeof(struct multiboot_mmap_entry);

    for (uint32_t i = 0; i < ctx.multi_mmap_size; i++) {
        struct multiboot_mmap_entry entry = ctx.multi_mmap[i];
        if (entry.type != MULTIBOOT_MEMORY_AVAILABLE) continue;

        uint32_t len  = (uint32_t)(entry.len_low | entry.len_high);
        uint32_t addr = (uint32_t)(entry.addr_low | entry.addr_high);
        printf("Region #%d: 0x%08X-0x%08X\n", i, addr, addr+len);
    }

    if (!mmap_init(&ctx)) {
        set_color(VGA_COLOR_RED);
        printf("Failed to create memory map\n");
        goto halt;
    }

    printf("Detected %d bytes of memory\n", ctx.available_bytes);

// Create and load an IDT
    for (int i = 0; i < 32; i++) {
        encode_idt_entry(g_idt, i, isr_stub_table[i], GDT_SEGMENT_SELECTOR(GDT_KERNEL_CODE, 0x00), TRAP_GATE_32 | INT_RING0 | INT_PRESENT);
    }

    g_idtr.size = (sizeof(struct idt_entry) * 256) - 1;
    g_idtr.base = (uint32_t) &g_idt;

    asm volatile ("lidt %0" :: "m"(g_idtr));

    struct rsdp *found_rsdp = acpi_locate_rsdp();
    if (!found_rsdp) {
        set_color(VGA_COLOR_RED);
        printf("Failed to find RSDP signature\n");
        goto halt;
    }

    if (found_rsdp->Revision == ACPI_VER_1) {
        if (!acpi_validate_rsdp_checksum(found_rsdp)) {
            set_color(VGA_COLOR_RED);
            printf("RSDP has an invalid checksum\n");
            goto halt;
        }
        ctx.rsdp = found_rsdp;
        ctx.xsdp = 0;

        if (!acpi_validate_sdt_checksum((struct ACPISDTHeader *)found_rsdp->RsdtAddress)) {
            set_color(VGA_COLOR_RED);
            printf("RSDT has an invalid checksum\n");
            goto halt;
        }

        ctx.rsdt = (struct rsdt*)found_rsdp->RsdtAddress;

    } else if (found_rsdp->Revision == ACPI_VER_OTHER) {
        set_color(VGA_COLOR_RED);
        printf("ACPI versions higher than 1.0 are not yet supported because I'm lazy\n");
        goto halt;
    } else {
        set_color(VGA_COLOR_RED);
        printf("Invalid RSDP\n");
        goto halt;
    }

    printf("Using ACPI v1.0\n");

    struct fadt *fadt = acpi_locate_sdt(ctx.rsdt, "FACP");
    if (!fadt) {
        set_color(VGA_COLOR_RED);
        printf("Failed to find FADT\n");
        goto halt;
    }
    printf("Found FADT at 0x%x\n", fadt);
    if (fadt->Flags & 1)
        printf("Legacy devices are supported\n");
    else {
        /*
        set_color(VGA_COLOR_RED);
        printf("Legacy devices are not supported. I'm too lazy to support modern devices, bye bye.\n");
        goto halt;
        */
    }
    pic_remap(PIC_1_START, PIC_2_START);
// ACPI version 1.0 is so old that we assume that our PC supports the 8042 ps/2 controller 
    initialize_8042ps2();
    encode_idt_entry(g_idt, 0x21, (uint32_t)&ps2_isr, GDT_SEGMENT_SELECTOR(GDT_KERNEL_CODE, 0x00), INT_GATE_32 | INT_RING0 | INT_PRESENT);
    asm volatile ("sti" ::);

    struct mcfg *mcfg = acpi_locate_sdt(ctx.rsdt, "MCFG");
    if (!mcfg) {
        set_color(VGA_COLOR_RED);
        printf("Failed to find MCFG\n");
    } else {

        printf("Looks like you are using PCIe- Found MCFG at 0x%x\n", mcfg);

        struct pci_config_space *ide = pcie_find_device(mcfg, MASS_STORAGE_CONTROLLER, IDE_INTERFACE);
        if (ide) { 
            printf("IDE controller detected. Program Interface: %X\n", ide->program_interface);

            uint16_t ide_buf[256] = {0};
            ide_identify(ATA_PRIMARY, ATA_MASTER, ide_buf);
        }
    }

    set_color(VGA_COLOR_WHITE);

    void *buf = mmap_find_first_free_block(&ctx);
    ata_read_sector(ATA_PRIMARY, ATA_MASTER, 0, (uint8_t *)buf);
//    struct fat12_boot_sector *boot_sect = (struct fat12_boot_sector *)buf;
//    int fat_size = boot_sect->total_sectors_16 * boot_sect->bytes_per_sector;

    write_tss(5, 0x10, 0x0);
    tss_flush(0x2B);
    void *stack_bottom = mmap_find_first_free_block(&ctx);
    void *stack_top = stack_bottom + 4096;
    switch_to_user_mode(stack_top, buf);
halt:
    for (;;) {}
}