/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   Interrupt Descriptor Table setup and dispatching
 * @license GPL-3.0-only
 */

#include <idt/idt.h>
#include <stdint.h>
#include <stddef.h>
#include <io/serial/serial.h>
#include <io/kbd/ps2.h>
#include <kernel.h>
#include <stdbool.h>
#include <sched/scheduler.h>
#include <config.h>
#include <sched/process.h>

struct interrupt_descriptor idt[256];
struct idtr idt_reg;

// Address to our first interrupt handler
extern char vector_0_handler[];

// Timer ticks
extern volatile uint64_t ticks;

/*
 * idt_set_entry - Sets an Interrupt Descriptor Table entry
 * @vector:  Vector number in the IDT
 * @handler: Pointer to the executable Interrupt Service Routine
 * @dpl:     Desired privilege level
 */
void idt_set_entry(uint8_t vector, void* handler, uint8_t dpl)
{
    uint64_t handler_addr = (uint64_t)handler;

    struct interrupt_descriptor* entry = &idt[vector];
    // Address is split in three parts so we right-shift progressively to get it all
    entry->address_low = handler_addr & 0xFFFF;
    entry->address_mid = (handler_addr >> 16) & 0xFFFF;
    entry->address_high = handler_addr >> 32;

    // Kernel code selector (as set in GDT)
    entry->selector = 0x8;
    // Interrupt gate, present, DPL (having: max DPL = 3)
    entry->flags = 0b1110 | ((dpl & 0b11) << 5) | (1 << 7);
    // We won't use IST for now
    entry->ist = 0;
}

/*
 * idt_load - Loads the Interrupt Descriptor Table
 * @idt_addr: Address to the IDT
 */
void idt_load(void* idt_addr)
{
    // "limit" = "size" = Size of the IDT - 1 byte = (16*256)-1 = 0xFFF
    idt_reg.limit = 0xFFF;
    idt_reg.base = (uint64_t)idt_addr;
    asm volatile("lidt %0" :: "m"(idt_reg));
}

/*
 * idt_init - Initializes the Interrupt Descriptor Table
 *
 * Sets all IDT entries and their corresponding service routines,
 * then loads it.
 */
void idt_init()
{
    for (size_t i=0; i<=KERNEL_IDT_ENTRIES; i++) {
        // Each vector handler is 16-byte aligned, so <vector_no>*16 = address of that handler
        idt_set_entry(i, vector_0_handler + (i*16), 0);
    }
    idt_load(&idt);
    DEBUG("IDT initialized");
}

/*
 * read_cr2 - Reads the CR2 register
 *
 * This function is useful because it gets the address
 * that the CPU tried to access in the case of a #PF.
 *
 * Return:
 * %val - CR2 register value
 */
static inline uint64_t read_cr2(void)
{
    uint64_t val;
    asm volatile ("mov %%cr2, %0" : "=r"(val));
    return val;
}

/*
 * page_fault_handler - Handler for #PF
 * @ctx: CPU context
 *
 * Shows detail about a #PF, especially what instruction (RIP)
 * caused it, and what address access (CR2) caused it.
 * Also displays an interpretation of the thrown error code.
 * Then halts the system. We could implement demand paging later.
 */
static void page_fault_handler(struct cpu_status_t* ctx)
{
    // It could be used to remap pages etc. to fix the fault, but right now what I'm more
    // interested in is getting more info out of those numbers cause i'm lost each time i have 
    // to read all this mess
    uint64_t cr2 = read_cr2();
    
    DEBUG("\x1b[38;5;231mPage Fault at rip=0x%p, err=%u (%s%s%s%s%s%s%s%s) when accessing addr=0x%p\x1b[0m", ctx->iret_rip, ctx->error_code,
    CHECK_BIT(ctx->error_code, 0) ? "PAGE_PROTECTION_VIOLATION " : "PAGE_NOT_PRESENT ",
    CHECK_BIT(ctx->error_code, 1) ? "ON_WRITE " : "ON_READ ",
    CHECK_BIT(ctx->error_code, 2) ? "IN_USER_MODE" : "IN_KERNEL_MODE",
    CHECK_BIT(ctx->error_code, 3) ? " WAS_RESERVED" : "",
    CHECK_BIT(ctx->error_code, 4) ? " ON_INSTRUCTION_FETCH" : "",
    CHECK_BIT(ctx->error_code, 5) ? " PK_VIOLATION" : "",
    CHECK_BIT(ctx->error_code, 6) ? " ON_SHADOWSTACK_ACCESS" : "",
    CHECK_BIT(ctx->error_code, 7) ? " SGX_VIOLATION" : "",
    cr2);

    panic(ctx, "page fault");
}

/*
 * gp_fault_handler - Handler for #GP
 * @ctx: CPU context
 *
 * Shows detail about a General Protection Fault,
 * and what may have caused it. Halts the system.
 */
static void gp_fault_handler(struct cpu_status_t* ctx)
{
    DEBUG("\x1b[38;5;231mGeneral Protection Fault at rip=0x%p, err=%u (%s)\x1b[0m",
    ctx->iret_rip,
    ctx->error_code,
    (ctx->error_code == 0) ? "NOT_SEGMENT_RELATED" : "SEGMENT_RELATED");

    // Segment-related
    if (ctx->error_code != 0) {
        bool is_external = CHECK_BIT(ctx->error_code, 0);
        // is it IDT, GDT, LDT?
        uint8_t table = ctx->error_code & 0x6; // 0b110 (isolate table)
        uint16_t index = ctx->error_code & 0xFFF8; // 13*1 1111111111111 + 000 = 1111111111111000
    
        char* table_names[4] = {"GDT", "IDT", "LDT", "IDT"};

        DEBUG("\x1b[38;5;231m%s in %s index %u\x1b[0m",
        is_external ? "EXTERNAL" : "INTERNAL",
        table_names[table],
        index);
    }

    panic(ctx, "gp fault");
}

// DEBUG
void kbdproc_main(void* arg)
{
    printf("Key pressed/released.\r\n");
}

/*
 * interrupt_dispatch - Interrupt dispatcher
 * @context: CPU context
 *
 * This function is where all interrupt routines go, after they passed
 * through their corresponding vector handler in the IDT assembly stub.
 * It catches all exceptions.
 *
 * Return:
 * <context> - CPU context after interrupt
 */
struct cpu_status_t* interrupt_dispatch(struct cpu_status_t* context)
{
    if (context == NULL) {
        panic(NULL, "Interrupt dispatch recieved NULL context!");
    }

    switch(context->vector_number) {
        case 0:
            panic(context, "Divide Error");
            break;
        case 1:
            panic(context, "Debug Exception");
            break;
        case 2:
            panic(context, "NMI Interrupt");
            break;
        case 3:
            panic(context, "Breakpoint Interrupt");
            break;
        case 4:
            panic(context, "Overflow Trap");
            break;
        case 5:
            panic(context, "BOUND Range Exceeded");
            break;
        case 6:
            panic(context, "Invalid Opcode");
            break;
        case 7:
            panic(context, "Device Not Available");
            break;
        case 8:
            panic(context, "Double Fault");
            break;
        case 9:
            panic(context, "Coprocessor Segment Overrun");
            break;
        case 10:
            panic(context, "Invalid TSS");
            break;
        case 11:
            panic(context, "Segment Not Present");
            break;
        case 12:
            panic(context, "Stack-Segment Fault");
            break;
        case 13:
            gp_fault_handler(context);
            break;
        case 14:
            page_fault_handler(context);
            break;
        case 15:
            panic(context, "Intel Reserved Interrupt (Achievement unlocked: How Did We Get Here?)");
            break;
        case 16:
            panic(context, "x87 Floating-Point Error");
            break;
        case 17:
            panic(context, "Alignment Check Fault");
            break;
        case 18:
            panic(context, "Machine Check");
            break;
        case 19:
            panic(context, "SIMD Floating-Point Exception");
            break;
        case 20:
            panic(context, "Virtualization Exception");
            break;
        case 21:
            panic(context, "Control Protection Exception");
            break;

        case 32: // Timer Interrupt
            ticks++;
            // Send an EOI so that we can continue having interrupts
            outb(0x20, 0x20);

            if (ticks % SCHEDULER_QUANTUM == 0) {
                return scheduler_schedule(context);
            }

            break;

        case 33: // Keyboard Interrupt
            keyboard_handler();
            process_create("keyboard-initiated", kbdproc_main, NULL); // DEBUG
            outb(0x20, 0x20);
            break;

        default:
            DEBUG("Unexpected Interrupt");
            break;
    }

    return context;
}