#include <stdint.h>

#include <lib/print.h>

#include <memory/mm.h>

#include "kbd.h"
#include "syscall.h"
#include "common.h"
#include "gdt.h"
#include "scheduler.h"
#include "process.h"

#include <drivers/pic/pic.h>
#include <drivers/ps2/ps2.h>
#include <lib/memory.h>

#define IDT_MAX_DESCRIPTORS 256
#define ATTR_KERN_32 0x8E // active(1), ring 0 (00), interrupt gate(0), 32 bit (1110)
#define ATTR_USER_32 0xEE // active, ring 3, 32bit
#define SEGMENT_KERN 0x08

typedef struct __attribute__((packed)) {
    uint16_t isr_low; 
    uint16_t selector; 
    uint8_t  reserved;  
    uint8_t  attributes;
    uint16_t isr_high; 
} idt_entry;

typedef struct __attribute__((packed)) {
    uint16_t limit;
    uint32_t base;
} idt_ptr;

// (interrupts.S)
extern void load_idt(void);
extern void isr0();
extern void isr1();
extern void isr8();
extern void isr13();
extern void isr14();
extern void isr32();
extern void isr33();

static void* isr_stub_table[256] = {
    [0] = isr0,
    [1] = isr1,
    [8] = isr8,
    [13] = isr13,
    [14] = isr14,
    [32] = isr32,
    [33] = isr33,
};

static const char* err_messages[32] = {
    [0] = "Divide by 0",
    [1] = "Debug",
    [4] = "Overflow",
    [6] = "Undefined Opcode",
    [8] = "",
    [13] = "General Protection Fault",
    [14] = "Page Fault",
};

idt_entry idt[IDT_MAX_DESCRIPTORS];
idt_ptr idtr;

static void set_idtr(uint8_t vector, void* isr, uint16_t sel, uint8_t attributes) {
    idt_entry* desc = &idt[vector];

    desc->isr_low    = (uint32_t)isr & 0xFFFF;
    desc->selector  = SEGMENT_KERN;
    desc->reserved   = 0;
    desc->attributes = attributes;
    desc->isr_high   = ((uint32_t)isr >> 16) & 0xFFFF;
}

void init_idt() {
    idtr.limit = (sizeof(idt_entry) * IDT_MAX_DESCRIPTORS) - 1;
    idtr.base  = (uint32_t)&idt;

    for (int i = 0; i < IDT_MAX_DESCRIPTORS; i++) {
        if (isr_stub_table[i] != 0) {
            set_idtr(i, isr_stub_table[i], SEGMENT_KERN, ATTR_KERN_32);
        } else {
            set_idtr(i, 0, SEGMENT_KERN, ATTR_KERN_32);
        }
    }

    set_idtr(0x80, syscall_handler, SEGMENT_KERN, ATTR_USER_32);

    load_idt();
}

typedef struct __attribute__((packed)) {
    uint32_t ds;
    uint32_t edi, esi, ebp, dummy, ebx, edx, ecx, eax;
    
    uint32_t int_no;    
    uint32_t error_code;
    
    uint32_t eip, cs, eflags, u_esp, u_ss;
} intr_regs_t;
void common_interrupt_handler(intr_regs_t *args) {
    switch (args->int_no) {
        case 0:
        case 1:
        case 2:
        case 3:
        case 4:
        case 5:
        case 6:
        case 7:
        case 8:
        case 9:
        case 10:
        case 11:
        case 12:
        case 13:
        case 15:
        case 16: {
            process_t* task = current_task();
            kprintf("\n--- CPU Exception ---\n");
            kprintf("Current Task PID: %d\n", task->pid);
            kprintf("Error: [%d] %s\n", args->int_no, err_messages[args->int_no]);
            kprintf("Error Code: 0x%x\n", args->error_code);
            kprintf("Instruction Pointer: 0x%x\n", args->eip);
            kprintf("Current Stack Pointer: 0x%x\n", (void*)args);
            kprintf("EAX: 0x%x\n", args->eax);
            kprintf("EFLAGS: 0x%x\n", args->eflags);
            kprintf("Code Segment: 0x%x\n", args->cs);
            kprintf("Calling Data Segment: 0x%x\n", args->ds);
            kprintf("Current cr3: 0x%x\n", fetch_cr3());
            kprintf("tss.esp0: 0x%x\n", tss.esp0);
            kprintf("instruction: 0x%x\n", *(uint32_t*)args->eip);
            kpanic();
            break;
        }
        case 14:{
            uint32_t faulting_address;
            asm volatile("mov %%cr2, %0" : "=r"(faulting_address));

            if (is_cow(faulting_address)) {
                uint32_t pde_entry = faulting_address >> 22;
                uint32_t pte_entry = (faulting_address >> 12) & 0x3FF;
                uint32_t* tablev = (uint32_t*)(CUR_PAGE_TABLES + (pde_entry * PAGE_SIZE));

                uint32_t p_dst = (uint32_t)p_alloc_frame();
                uint32_t p_src = (uint32_t)tablev[pte_entry] & 0xFFFFF000;

                map_page(TEMP_PT_SRC, p_src, PAGE_WRITABLE | PAGE_PRESENT);
                map_page(TEMP_PT_DST, p_dst, PAGE_WRITABLE | PAGE_PRESENT);

                memcpy((void*)TEMP_PT_SRC, (void*)TEMP_PT_DST, PAGE_SIZE);

                unmap_page(TEMP_PT_SRC);
                unmap_page(TEMP_PT_DST);

                uint32_t flags = tablev[pte_entry] & 0xFFF;
                flags &= ~PAGE_COW;
                tablev[pte_entry] = p_dst | flags | PAGE_WRITABLE;

                flush_tlb();
                break;
            }

            process_t* task = current_task();
            kprintf("\n--- PAGE FAULT ---\n");
            kprintf("Current Task PID: %d\n", task->pid);
            kprintf("Faulted Virtual Address: 0x%x\n", faulting_address);
            kprintf("Error Code: 0x%x\n", args->error_code);
            kprintf("Instruction Pointer: 0x%x\n", args->eip);
            kprintf("Current Stack Pointer: 0x%x\n", (void*)args);
            kprintf("EAX: 0x%x\n", args->eax);
            kprintf("EFLAGS: 0x%x\n", args->eflags);
            kprintf("Code Segment: 0x%x\n", args->cs);
            kprintf("Current cr3: 0x%x\n", fetch_cr3());
            if (args->cs == 0x1b) {
                kprintf("User ESP: 0x%x\n", args->u_esp);
                kprintf("ESP: 0x%x\t0x%x\t0x%x\t0x%x\n", 
                    *(uint32_t*)args->u_esp,
                    *(uint32_t*)(args->u_esp + 4),
                    *(uint32_t*)(args->u_esp + 8),
                    *(uint32_t*)(args->u_esp + 12));
                kprintf("User SS: 0x%x\n", args->u_ss);
            }
            kprintf("Caused by: %s, %s, running in %s mode\n",
                (args->error_code & 0x01) ? "Protection Violation" : "Page Not Present",
                (args->error_code & 0x02) ? "Write" : "Read",
                (args->error_code & 0x04) ? "User" : "Supervisor");
            kpanic();
            break;
        }
        case 32: {

            break;
        }
        case 33: {
            uint8_t data = ps2_read_data();
            char ascii = ps2_get_ascii(data);
            if (ascii != 0) {
                kbd_push(ascii);
            }
            send_eoi_master();
            break;
        }
        case 46: { // IDE PRIMARY
            kprintf("IDE PRIMARY interrupt\n");
            send_eoi_slave();
            break;
        }
        case 47: {  // IDE SLAVE
            kprintf("IDE SLAVE interrupt\n");
            send_eoi_slave();
            break;
        }
        default:
            kprintf("Received unhandled interrupt: %d\n", args->int_no);
            return;
    }
}