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

#define MAGIC_BREAKPOINT __asm__ volatile("xchgw %bx, %bx");
#define FD_STDOUT 0
#define FD_STDIN  1

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();
extern struct gdt_entry gdt[6];
extern void syscall_isr();
extern void syscall();
extern void test();

linked_mem_t stdout = {0};
linked_mem_t stdin  = {0};

struct tss_entry tss;

struct stack_frame {
    uint32_t edi, esi, ebp, esp, ebx, edx, ecx, eax;
};

void syscall_read(struct stack_frame regs) {
    uint32_t fd    = regs.edi;
    char *buf      = (char *)regs.esi;
    uint32_t count = regs.edx;

    if (fd == FD_STDOUT) {
        int pages = count / BLOCK_SIZE;
        if (pages == 0) {
            memcpy(buf, stdout.data, count);
        }
    }
}

void syscall_write(struct stack_frame regs) {
    uint32_t fd    = regs.edi;
    char *buf      = (char *)regs.esi;
    //uint32_t count = regs.edx;

    if (fd == FD_STDOUT) {
        printf("%s", buf);
    }
}

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 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);

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

    printf("[KERNEL] 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("[KERNEL] 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("[KERNEL] Found FADT at 0x%x\n", fadt);
    if (fadt->Flags & 1)
        printf("[KERNEL] 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 

    // Allocate 4KiB for stdout
    stdout.data = mmap_find_first_free_block(&ctx);
    stdout.next = NULL;
    // Allocate 4KiB for stdin
    stdin.data = mmap_find_first_free_block(&ctx);
    stdin.next = NULL;

    initialize_8042ps2();
    encode_idt_entry(g_idt, 0x21, (uint32_t)&ps2_isr, GDT_SEGMENT_SELECTOR(GDT_KERNEL_CODE, RING0), 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);
        }
    }
*/

    write_tss(&tss, gdt, 5, 0x10, 0x0);
    tss_flush(0x2B);
    encode_idt_entry(g_idt, 0x80, (uint32_t)&syscall_isr, GDT_SEGMENT_SELECTOR(GDT_USER_CODE, RING3), INT_GATE_32 | INT_RING3 | INT_PRESENT);

    uint16_t cs_selector = 0x10;  // Example code segment selector
    uint32_t eip = (uint32_t)&syscall;    // Example entry point address
    uint32_t esp;

    // Get the current value of ESP
    __asm__ volatile ("mov %%esp, %0" : "=r"(esp));

    // Write to MSR_SYSENTER_CS (0x174)
    __asm__ volatile (
        "movl $0x174, %%ecx \n\t"   // MSR address
        "movl %0, %%eax \n\t"       // Lower 32 bits (cs_selector)
        "xor %%edx, %%edx \n\t"     // Upper 32 bits (zero for 32-bit value)
        "wrmsr"
        : 
        : "r"((uint32_t)cs_selector)
        : "eax", "ecx", "edx"
    );

    // Write to MSR_SYSENTER_ESP (0x175)
    __asm__ volatile (
        "movl $0x175, %%ecx \n\t"   // MSR address
        "movl %0, %%eax \n\t"       // Lower 32 bits (esp)
        "xor %%edx, %%edx \n\t"     // Upper 32 bits (zero for 32-bit value)
        "wrmsr"
        : 
        : "r"(esp)
        : "eax", "ecx", "edx"
    );

    // Write to MSR_SYSENTER_EIP (0x176)
    __asm__ volatile (
        "movl $0x176, %%ecx \n\t"   // MSR address
        "movl %0, %%eax \n\t"       // Lower 32 bits (eip)
        "xor %%edx, %%edx \n\t"     // Upper 32 bits (zero for 32-bit value)
        "wrmsr"
        : 
        : "r"(eip)
        : "eax", "ecx", "edx"
    );

    void *page = mmap_find_first_free_block(&ctx);
    ata_read_sector(ATA_PRIMARY, ATA_MASTER, 0, page);
    printf("[KERNEL] Init program loaded to 0x%08X\n", page);
    printf("[KERNEL] Transferring control... goodbye\n");
    set_color(VGA_COLOR_WHITE);
    // 4 KiB for stack
    void *stack_bottom = mmap_find_first_free_block(&ctx);
    void *stack_top = stack_bottom + 0x1000;
    MAGIC_BREAKPOINT
    switch_to_user_mode(stack_top, page);
halt:
    for (;;) {}
}