#include "../console.h"
#include "../drivers/keyboard.h"
#include "../drivers/pit.h"
#include "../drivers/port.h"
#include "../drivers/vga.h"
#include "../proc.h"
#include "../syscall.h"
#include "gdt.h"
#include "isr.h"
#include "memlayout.h"

enum {
    IDT_HANDLERS = 256,
};

typedef struct {
    uint16_t low_offset;
    uint16_t selector;
    uint8_t always0;
    uint8_t type : 4;
    uint8_t s : 1;
    uint8_t dpl : 2;
    uint8_t p : 1;
    uint16_t high_offset;
} __attribute__((packed)) idt_gate_t;

idt_gate_t idt[IDT_HANDLERS];

#define low_16(address) (uint16_t)((address) & 0xFFFF)
#define high_16(address) (uint16_t)(((address) >> 16) & 0xFFFF)

#define STS_IG32 0xE // 32-bit Interrupt Gate
#define STS_TG32 0xF // 32-bit Trap Gate

void set_idt_gate(int n, _Bool istrap, uint32_t handler, uint8_t dpl) {
    idt[n].low_offset = low_16(handler);
    idt[n].selector = 0x08; // see GDT
    idt[n].always0 = 0;
    idt[n].type = istrap ? STS_TG32 : STS_IG32;
    idt[n].s = 0;
    idt[n].dpl = dpl;
    idt[n].p = 1;
    idt[n].high_offset = high_16(handler);
}

// defined in vectors.S
extern const uint32_t default_handlers[];

void init_idt() {
    if (default_handlers[0] == 0) {
        panic("handler table empty\n");
    }
    for (int i = 0; i < IDT_HANDLERS; i++) {
        set_idt_gate(i, 0, default_handlers[i], 0);
    }
    set_idt_gate(T_SYSCALL, 1, default_handlers[T_SYSCALL], DPL_USER);
}

const char *const exception_messages[] = {
    [0] = "Division By Zero",
    [1] = "Debug",
    [2] = "Non Maskable Interrupt",
    [3] = "Breakpoint",
    [4] = "Into Detected Overflow",
    [5] = "Out of Bounds",
    [6] = "Invalid Opcode",
    [7] = "No Coprocessor",

    [8] = "Double Fault",
    [9] = "Coprocessor Segment Overrun",
    [10] = "Bad TSS",
    [11] = "Segment Not Present",
    [12] = "Stack Fault",
    [13] = "General Protection Fault",
    [14] = "Page Fault",
    [15] = "Unknown Interrupt",

    [16] = "Coprocessor Fault",
    [17] = "Alignment Check",
    [18] = "Machine Check",
};

#define ARRLEN(a) (sizeof(a) / sizeof(a[0]))

static isr_t interrupt_handlers[IDT_HANDLERS];

void register_interrupt_handler(uint8_t i, isr_t handler) {
    interrupt_handlers[i] = handler;
}

void trap(registers_t *r) {
    // EOI
    if (r->int_no >= 40) {
        port_byte_out(0xA0, 0x20); /* follower */
    }
    if (r->int_no >= 32) {
        port_byte_out(0x20, 0x20); /* leader */
    }

    // Call registered handler
    if (interrupt_handlers[r->int_no] != 0) {
        isr_t handler = interrupt_handlers[r->int_no];
        handler(r);
        return;
    }

    if (r->int_no < 32) {
        const char *msg = "Reserved";
        if (r->int_no < ARRLEN(exception_messages)) {
            msg = exception_messages[r->int_no];
        }
        if (r->cs & 3) {
            // exception from user mode, kill offending process
            printk("Exception: ");
            printk(msg);
            printk("\n");
            killproc();
        }
        panic(msg);
    }
}

/* takes a ptr that is supposed to be from userspace. If it is not, we
 * return 0, if it is a valid we return 1
 */
bool is_userspace_ptr_mapped(uint32_t ptr) {
    if (ptr >= KERNBASE) {
        return 0;
    }

    pde_t *pgdir = get_user_proc_page_directory();
    if (!pgdir) {
        return 0;
    }

    pde_t pde = pgdir[PDX(ptr)];
    if ((pde & (PTE_P | PTE_U)) != (PTE_P | PTE_U)) {
        return 0;
    }
    if (pde & PDE_PS) {
        return 0;
    }

    pte_t *table = (pte_t *)P2V(PTE_ADDR(pde));
    pte_t pte = table[PTX(ptr)];
    if ((pte & (PTE_P | PTE_U)) != (PTE_P | PTE_U)) {
        return 0;
    }

    return 1;
}

static bool is_userspace_range_mapped(uint32_t ptr, uint32_t size) {
    if (size == 0) {
        return 1;
    }
    if (ptr >= KERNBASE) {
        return 0;
    }

    uint32_t end = ptr + size - 1;
    if (end < ptr || end >= KERNBASE) {
        return 0;
    }

    uint32_t last_page = PGROUNDDOWN(end);
    for (uint32_t addr = ptr;; addr = PGROUNDDOWN(addr) + PGSIZE) {
        if (!is_userspace_ptr_mapped(addr)) {
            return 0;
        }
        if (PGROUNDDOWN(addr) == last_page) {
            return 1;
        }
    }
}

static bool is_userspace_cstr(uint32_t ptr) {
    for (uint32_t addr = ptr;; addr++) {
        if (addr == 0 || !is_userspace_ptr_mapped(addr)) {
            return 0;
        }
        if (*(const char *)addr == '\0') {
            return (void *)ptr;
        }
    }
}

static _Noreturn void userspace_panic(const char *msg) {
    if (!vga_is_text_mode()) {
        switch_to_text_mode();
        vga_clear_screen();
    }
    printk(msg);
    killproc();
}

static void handle_puts(uintptr_t s) {
    if (!is_userspace_cstr(s)) {
        userspace_panic("SYS_puts panic: page fault\n");
    }
    printk((const char *)s);
}

static void require_text_mode_for_userspace_text_syscall() {
    if (!vga_is_text_mode()) {
        userspace_panic("Userspace panic: text syscall in graphics mode\n");
    }
}

static void handle_swap_frame(uintptr_t frame) {
    enum {
        VGA_GRAPHICS_FRAME_SIZE = VGA_GRAPHICS_WIDTH * VGA_GRAPHICS_HEIGHT,
    };

    if (vga_is_text_mode()) {
        userspace_panic("Userspace panic: frame swap in text mode\n");
    }
    if (!is_userspace_range_mapped(frame, VGA_GRAPHICS_FRAME_SIZE)) {
        userspace_panic("SYS_swap_frame panic: page fault\n");
    }

    uint8_t *video = (uint8_t *)(KERNBASE + 0xA0000);
    uint8_t *user = (uint8_t *)frame;
    for (uint32_t i = 0; i < VGA_GRAPHICS_FRAME_SIZE; i++) {
        video[i] = user[i];
    }
}

static void refill_keyboard_copy_buffer(void) {
    size_t count = kbd_state_shrd.len;
    if (count == 0) {
        return;
    }

    cli();
    size_t rem = KEYBOARD_INTERRUPT_BUF_CAP - kbd_state_shrd.copy_len;
    size_t copying = rem < count ? rem : count;
    memcpy(kbd_state_shrd.copy_buf, kbd_state_shrd.buf, copying);
    kbd_state_shrd.len -= copying;
    kbd_state_shrd.copy_len += copying;
    sti();
}

static int handle_getc(void) {
    if (kbd_can_take_from_copy_buffer()) {
        return kbd_take_from_copy_buffer();
    }

    refill_keyboard_copy_buffer();
    if (!kbd_can_take_from_copy_buffer()) {
        return -1;
    }
    return kbd_take_from_copy_buffer();
}

static void handle_syscall(registers_t *r) {
    switch (r->eax) {
    case SYS_exit:
        if (r->ebx == 0) {
            printk("* success\n");
        } else {
            printk("* failure\n");
        }
        killproc();
    case SYS_greet:
        require_text_mode_for_userspace_text_syscall();
        printk("Hello world!\n");
        r->eax = 0;
        break;
    case SYS_putc:
        require_text_mode_for_userspace_text_syscall();
        printk((const char[]){(char)r->ebx, '\0'});
        r->eax = 0;
        break;
    case SYS_puts:
        require_text_mode_for_userspace_text_syscall();
        handle_puts(r->ebx);
        r->eax = 0;
        break;
    case SYS_switch_to_text:
        switch_to_text_mode();
        r->eax = 0;
        break;
    case SYS_switch_to_graphics:
        switch_to_graphics_mode();
        r->eax = 0;
        break;
    case SYS_swap_frame:
        handle_swap_frame(r->ebx);
        r->eax = 0;
        break;
    case SYS_time_ms:
        r->eax = get_uptime_ms();
        break;
    case SYS_halt:
        asm volatile("hlt");
        r->eax = 0;
        break;
    case SYS_getc:
        r->eax = handle_getc();
        break;
    case SYS_set_beep:
        if (r->ebx > MAX_BEEP_FREQUENCY_HZ) {
            userspace_panic("Userspace panic: beep frequency out of range\n");
        }
        set_beep_frequency_hz(r->ebx);
        r->eax = 0;
        break;
    default:
        userspace_panic("Userspace panic: Unknown syscall\n");
    }
}

static void init_pic() {
    // ICW1
    port_byte_out(0x20, 0x11);
    port_byte_out(0xA0, 0x11);

    // ICW2
    port_byte_out(0x21, 0x20);
    port_byte_out(0xA1, 0x28);

    // ICW3
    port_byte_out(0x21, 0x04);
    port_byte_out(0xA1, 0x02);

    // ICW4
    port_byte_out(0x21, 0x01);
    port_byte_out(0xA1, 0x01);

    // OCW1
    port_byte_out(0x21, 0x0);
    port_byte_out(0xA1, 0x0);
}

typedef struct {
    uint16_t limit;
    void *base;
} __attribute__((packed)) idt_register_t;

static idt_register_t idt_reg;

void load_idt() {
    init_idt();

    idt_reg.base = &idt;
    idt_reg.limit = sizeof(idt) - 1;
    asm("lidt (%0)" : : "r"(&idt_reg));

    init_pic();

    register_interrupt_handler(T_SYSCALL, handle_syscall);
}

void cli() { asm("cli"); }

void sti() { asm("sti"); }