#include "idt.h"

#define IDT_SIZE 256
#define IDT_TYPE_INTERRUPT 0x8E

typedef struct {
    UINT16 offset_low;
    UINT16 selector;
    UINT8 ist;
    UINT8 type_attr;
    UINT16 offset_mid;
    UINT32 offset_high;
    UINT32 zero;
} __attribute__((packed)) IdtEntry;

typedef struct {
    UINT16 limit;
    UINT64 base;
} __attribute__((packed)) IdtPtr;

static IdtEntry idt[IDT_SIZE];
static BootInfo *gBoot = NULL;

extern void (*isr_stub_table[])(void);

static void idt_set_gate(UINTN index, void (*handler)(void))
{
    UINT64 addr = (UINT64)(UINTN)handler;
    UINT16 selector = 0;

    __asm__ __volatile__("mov %%cs, %0" : "=r"(selector));

    idt[index].offset_low = (UINT16)(addr & 0xFFFF);
    idt[index].selector = selector;
    idt[index].ist = 0;
    idt[index].type_attr = IDT_TYPE_INTERRUPT;
    idt[index].offset_mid = (UINT16)((addr >> 16) & 0xFFFF);
    idt[index].offset_high = (UINT32)((addr >> 32) & 0xFFFFFFFF);
    idt[index].zero = 0;
}

static void lidt(const IdtPtr *idtr)
{
    __asm__ __volatile__("lidt (%0)" :: "r"(idtr));
}

static const CHAR16 *exception_name(UINTN vector)
{
    switch (vector) {
    case 0: return L"Divide Error";
    case 1: return L"Debug";
    case 2: return L"Non-Maskable Interrupt";
    case 3: return L"Breakpoint";
    case 4: return L"Overflow";
    case 5: return L"Bound Range Exceeded";
    case 6: return L"Invalid Opcode";
    case 7: return L"Device Not Available";
    case 8: return L"Double Fault";
    case 9: return L"Coprocessor Segment Overrun";
    case 10: return L"Invalid TSS";
    case 11: return L"Segment Not Present";
    case 12: return L"Stack-Segment Fault";
    case 13: return L"General Protection Fault";
    case 14: return L"Page Fault";
    case 15: return L"Reserved";
    case 16: return L"x87 Floating-Point";
    case 17: return L"Alignment Check";
    case 18: return L"Machine Check";
    case 19: return L"SIMD Floating-Point";
    case 20: return L"Virtualization";
    case 21: return L"Control Protection";
    case 22: return L"Reserved";
    case 23: return L"Reserved";
    case 24: return L"Reserved";
    case 25: return L"Reserved";
    case 26: return L"Reserved";
    case 27: return L"Reserved";
    case 28: return L"Hypervisor Injection";
    case 29: return L"VMM Communication";
    case 30: return L"Security";
    case 31: return L"Reserved";
    default: return L"Unknown";
    }
}

static inline void outb(UINT16 port, UINT8 value)
{
    __asm__ __volatile__("outb %0, %1" :: "a"(value), "Nd"(port));
}

static void pic_eoi(UINTN vector)
{
    if (vector >= 40) {
        outb(0xA0, 0x20);  /* EOI to slave PIC */
    }
    outb(0x20, 0x20);      /* EOI to master PIC */
}

static void halt_forever(void)
{
    for (;;) {
        __asm__ __volatile__("cli; hlt");
    }
}

void isr_handler(ISRFrame *frame)
{
    UINT64 cr2 = 0;

    /* Hardware IRQs (vectors 32-47): send EOI and return */
    if (frame->vector >= 32 && frame->vector <= 47) {
        pic_eoi(frame->vector);
        return;
    }

    /* CPU exceptions (vectors 0-31): print diagnostics and halt */
    if (gBoot != NULL && gBoot->print != NULL) {
        gBoot->print(L"\n\rEXCEPTION: %d (%s)\n\r", frame->vector,
                     exception_name(frame->vector));
        gBoot->print(L"  Error Code: 0x%lx\n\r", frame->error_code);
        gBoot->print(L"  RIP: 0x%lx  CS: 0x%lx  RFLAGS: 0x%lx\n\r",
                     frame->rip, frame->cs, frame->rflags);
    }

    if (frame->vector == 14) {
        __asm__ __volatile__("mov %%cr2, %0" : "=r"(cr2));
        if (gBoot != NULL && gBoot->print != NULL) {
            gBoot->print(L"  CR2: 0x%lx\n\r", cr2);
        }
    }

    halt_forever();
}

void idt_init(BootInfo *Boot)
{
    IdtPtr old_idtr;
    IdtPtr idtr;
    IdtEntry *old_idt = NULL;
    UINTN i = 0;

    gBoot = Boot;

    /* Read the firmware's existing IDT so we can preserve its entries */
    __asm__ __volatile__("sidt %0" : "=m"(old_idtr));
    old_idt = (IdtEntry *)(UINTN)old_idtr.base;

    /* Copy the entire existing IDT first (preserves firmware IRQ handlers) */
    for (i = 0; i < IDT_SIZE; i++) {
        if (old_idt != NULL && (i * sizeof(IdtEntry)) < (UINTN)(old_idtr.limit + 1)) {
            idt[i] = old_idt[i];
        } else {
            idt_set_gate(i, isr_stub_table[i]);
        }
    }

    /* Override only CPU exception vectors (0-31) with our handlers */
    for (i = 0; i < 32; i++) {
        idt_set_gate(i, isr_stub_table[i]);
    }

    idtr.limit = (UINT16)(sizeof(idt) - 1);
    idtr.base = (UINT64)(UINTN)idt;

    lidt(&idtr);
}