#include <stdint.h>

#include "apic.h"
// for debugging
#include "print.h"
#include "kbd.h"
#include "common.h"

// for trap debugging
u64 dbg_var = 0;

pstruct int_desc {
  u16 addr_low;
  u16 selector;
  u8 ist;
  u8 flags;
  u16 addr_mid;
  u32 addr_high;
  u32 resv;
};

enum INT_DESC_FLAGS {
  ID_GATE_INT = 0b00001110,
  ID_GATE_TRAP = 0b00001111,
  ID_DPL = 0b00100000,
  ID_PRESENT = 0b10000000,
};

#define CODE64_SELECTOR 0x28

struct int_desc idt[256];

void set_idt(u8 vect, void *handler, u8 dpl) {
  struct int_desc *ent = &idt[vect];
  u64 addr = (u64)handler;
  ent->addr_low = addr & 0xFFFF;
  ent->addr_mid = (addr >> 16) & 0xFFFF;
  ent->addr_high = addr >> 32;

  ent->selector = CODE64_SELECTOR;
  ent->flags = ID_GATE_INT | dpl * ID_DPL | ID_PRESENT;
  ent->ist = 0;
}

pstruct idtr {
  u16 limit;
  u64 base;
};

#define IDTR_LIMIT (sizeof(struct int_desc) * 256 - 1)

void load_idt(void *addr) {
  struct idtr ir = (struct idtr){
      .limit = IDTR_LIMIT,
      .base = (u64)addr,
  };

  asm volatile("lidt %0" ::"m"(ir));
}

extern void draw(u8 vect);

struct cpu_ctx {
  u64 r15, r14, r13, r12, r11, r10, r9, r8, rdi, rsi, rdx, rcx, rbx, rax;

  u64 vect;
  u64 err;

  u64 iret_rip;
  u64 iret_cs;
  u64 iret_flags;
  u64 iret_rsp;
  u64 iret_ss;
};

struct cpu_ctx *interrupt_dispatch(struct cpu_ctx *ctx) {
  if(ctx->vect = 0x21) {
    kbd_handler();
    return ctx;
  }
  print("\n[IRQ ");
  print8(ctx->vect);
  print("] with message: ");
  print64(dbg_var);
  print("\n");

  return ctx;
}

__attribute__((naked)) void int_stub(void) {
  __asm__ volatile("push %%rax\n"
                   "push %%rbx\n"
                   "push %%rcx\n"
                   "push %%rdx\n"
                   "push %%rsi\n"
                   "push %%rdi\n"
                   "push %%r8\n"
                   "push %%r9\n"
                   "push %%r10\n"
                   "push %%r11\n"
                   "push %%r12\n"
                   "push %%r13\n"
                   "push %%r14\n"
                   "push %%r15\n"
                   "mov %%rsp, %%rdi\n"
                   "call interrupt_dispatch\n"
                   "mov %%rax, %%rsp\n"
                   "pop %%r15\n"
                   "pop %%r14\n"
                   "pop %%r13\n"
                   "pop %%r12\n"
                   "pop %%r11\n"
                   "pop %%r10\n"
                   "pop %%r9\n"
                   "pop %%r8\n"
                   "pop %%rdi\n"
                   "pop %%rsi\n"
                   "pop %%rdx\n"
                   "pop %%rcx\n"
                   "pop %%rbx\n"
                   "pop %%rax\n"
                   "add $16, %%rsp\n"
                   "iretq\n"
                   :
                   :
                   : "memory");
}

// Macro to generate interrupt stubs without error codes (push dummy 0)
#define ISR_NO_ERR(num)                                                        \
  __attribute__((naked)) void isr##num(void) {                                 \
    asm volatile("push $0\n"                                                   \
                 "push %0\n"                                                   \
                 "jmp int_stub\n"                                              \
                 :                                                             \
                 : "i"(num));                                                  \
  }

// Macro to generate interrupt stubs with error codes
#define ISR_ERR(num)                                                           \
  __attribute__((naked)) void isr##num(void) {                                 \
    asm volatile("push %0\n"                                                   \
                 "jmp int_stub\n"                                              \
                 :                                                             \
                 : "i"(num));                                                  \
  }

// CPU exceptions (0-31)
ISR_NO_ERR(0)  // Divide by zero
ISR_NO_ERR(1)  // Debug
ISR_NO_ERR(2)  // NMI
ISR_NO_ERR(3)  // Breakpoint
ISR_NO_ERR(4)  // Overflow
ISR_NO_ERR(5)  // Bound range exceeded
ISR_NO_ERR(6)  // Invalid opcode
ISR_NO_ERR(7)  // Device not available
ISR_ERR(8)     // Double fault (pushes error code)
ISR_NO_ERR(9)  // Coprocessor segment overrun
ISR_ERR(10)    // Invalid TSS (pushes error code)
ISR_ERR(11)    // Segment not present (pushes error code)
ISR_ERR(12)    // Stack-segment fault (pushes error code)
ISR_ERR(13)    // General protection fault (pushes error code)
ISR_ERR(14)    // Page fault (pushes error code)
ISR_NO_ERR(15) // Reserved
ISR_NO_ERR(16) // x87 FPU error
ISR_ERR(17)    // Alignment check (pushes error code)
ISR_NO_ERR(18) // Machine check
ISR_NO_ERR(19) // SIMD floating-point exception
ISR_NO_ERR(20) // Virtualization exception
ISR_ERR(21)    // Control protection exception (pushes error code)
ISR_NO_ERR(22) // Reserved
ISR_NO_ERR(23) // Reserved
ISR_NO_ERR(24) // Reserved
ISR_NO_ERR(25) // Reserved
ISR_NO_ERR(26) // Reserved
ISR_NO_ERR(27) // Reserved
ISR_NO_ERR(28) // Reserved
ISR_NO_ERR(29) // Reserved
ISR_ERR(30)    // Security exception (pushes error code)
ISR_NO_ERR(31) // Reserved

// IRQs and user-defined interrupts (32-255)
ISR_NO_ERR(32)
ISR_NO_ERR(33) ISR_NO_ERR(34) ISR_NO_ERR(35) ISR_NO_ERR(36) ISR_NO_ERR(37) ISR_NO_ERR(
    38) ISR_NO_ERR(39) ISR_NO_ERR(40) ISR_NO_ERR(41) ISR_NO_ERR(42) ISR_NO_ERR(43)
    ISR_NO_ERR(44) ISR_NO_ERR(45) ISR_NO_ERR(46) ISR_NO_ERR(47) ISR_NO_ERR(
        48) ISR_NO_ERR(49) ISR_NO_ERR(50) ISR_NO_ERR(51) ISR_NO_ERR(52) ISR_NO_ERR(53)
        ISR_NO_ERR(54) ISR_NO_ERR(55) ISR_NO_ERR(56) ISR_NO_ERR(57) ISR_NO_ERR(
            58) ISR_NO_ERR(59) ISR_NO_ERR(60) ISR_NO_ERR(61) ISR_NO_ERR(62)
            ISR_NO_ERR(63) ISR_NO_ERR(64) ISR_NO_ERR(65) ISR_NO_ERR(66) ISR_NO_ERR(
                67) ISR_NO_ERR(68) ISR_NO_ERR(69) ISR_NO_ERR(70) ISR_NO_ERR(71)
                ISR_NO_ERR(72) ISR_NO_ERR(73) ISR_NO_ERR(74) ISR_NO_ERR(75) ISR_NO_ERR(
                    76) ISR_NO_ERR(77) ISR_NO_ERR(78) ISR_NO_ERR(79) ISR_NO_ERR(80)
                    ISR_NO_ERR(81) ISR_NO_ERR(82) ISR_NO_ERR(83) ISR_NO_ERR(
                        84) ISR_NO_ERR(85) ISR_NO_ERR(86) ISR_NO_ERR(87)
                        ISR_NO_ERR(88) ISR_NO_ERR(89) ISR_NO_ERR(90) ISR_NO_ERR(
                            91) ISR_NO_ERR(92) ISR_NO_ERR(93) ISR_NO_ERR(94)
                            ISR_NO_ERR(95) ISR_NO_ERR(96) ISR_NO_ERR(97) ISR_NO_ERR(
                                98) ISR_NO_ERR(99) ISR_NO_ERR(100) ISR_NO_ERR(101)
                                ISR_NO_ERR(102) ISR_NO_ERR(103) ISR_NO_ERR(
                                    104) ISR_NO_ERR(105) ISR_NO_ERR(106)
                                    ISR_NO_ERR(107) ISR_NO_ERR(108) ISR_NO_ERR(
                                        109) ISR_NO_ERR(110) ISR_NO_ERR(111)
                                        ISR_NO_ERR(112) ISR_NO_ERR(113) ISR_NO_ERR(
                                            114) ISR_NO_ERR(115) ISR_NO_ERR(116)
                                            ISR_NO_ERR(117) ISR_NO_ERR(118) ISR_NO_ERR(
                                                119) ISR_NO_ERR(120) ISR_NO_ERR(121)
                                                ISR_NO_ERR(122) ISR_NO_ERR(
                                                    123) ISR_NO_ERR(124)
                                                    ISR_NO_ERR(125) ISR_NO_ERR(
                                                        126) ISR_NO_ERR(127)
                                                        ISR_NO_ERR(128) ISR_NO_ERR(
                                                            129) ISR_NO_ERR(130)
                                                            ISR_NO_ERR(131) ISR_NO_ERR(
                                                                132) ISR_NO_ERR(133)
                                                                ISR_NO_ERR(
                                                                    134) ISR_NO_ERR(135)
                                                                    ISR_NO_ERR(136) ISR_NO_ERR(
                                                                        137) ISR_NO_ERR(138)
                                                                        ISR_NO_ERR(139) ISR_NO_ERR(
                                                                            140) ISR_NO_ERR(141)
                                                                            ISR_NO_ERR(142) ISR_NO_ERR(
                                                                                143)
                                                                                ISR_NO_ERR(144) ISR_NO_ERR(145) ISR_NO_ERR(146) ISR_NO_ERR(147) ISR_NO_ERR(148) ISR_NO_ERR(149) ISR_NO_ERR(150) ISR_NO_ERR(151) ISR_NO_ERR(152) ISR_NO_ERR(153) ISR_NO_ERR(154) ISR_NO_ERR(155) ISR_NO_ERR(156) ISR_NO_ERR(157) ISR_NO_ERR(158) ISR_NO_ERR(159) ISR_NO_ERR(160) ISR_NO_ERR(161) ISR_NO_ERR(162) ISR_NO_ERR(163) ISR_NO_ERR(164) ISR_NO_ERR(165) ISR_NO_ERR(166) ISR_NO_ERR(167) ISR_NO_ERR(168) ISR_NO_ERR(169) ISR_NO_ERR(170) ISR_NO_ERR(171) ISR_NO_ERR(172) ISR_NO_ERR(173) ISR_NO_ERR(
                                                                                    174) ISR_NO_ERR(175)
                                                                                    ISR_NO_ERR(
                                                                                        176) ISR_NO_ERR(177)
                                                                                        ISR_NO_ERR(
                                                                                            178) ISR_NO_ERR(179)
                                                                                            ISR_NO_ERR(
                                                                                                180) ISR_NO_ERR(181)
                                                                                                ISR_NO_ERR(
                                                                                                    182) ISR_NO_ERR(183)
                                                                                                    ISR_NO_ERR(
                                                                                                        184) ISR_NO_ERR(185)
                                                                                                        ISR_NO_ERR(
                                                                                                            186)
                                                                                                            ISR_NO_ERR(
                                                                                                                187)
                                                                                                                ISR_NO_ERR(188) ISR_NO_ERR(189) ISR_NO_ERR(190) ISR_NO_ERR(191) ISR_NO_ERR(192) ISR_NO_ERR(193) ISR_NO_ERR(194) ISR_NO_ERR(195) ISR_NO_ERR(
                                                                                                                    196) ISR_NO_ERR(197)
                                                                                                                    ISR_NO_ERR(198) ISR_NO_ERR(
                                                                                                                        199)
                                                                                                                        ISR_NO_ERR(
                                                                                                                            200) ISR_NO_ERR(201) ISR_NO_ERR(202) ISR_NO_ERR(203) ISR_NO_ERR(204) ISR_NO_ERR(205) ISR_NO_ERR(206) ISR_NO_ERR(207)
                                                                                                                            ISR_NO_ERR(208) ISR_NO_ERR(209) ISR_NO_ERR(210) ISR_NO_ERR(211) ISR_NO_ERR(212) ISR_NO_ERR(213) ISR_NO_ERR(214) ISR_NO_ERR(
                                                                                                                                215) ISR_NO_ERR(216) ISR_NO_ERR(217) ISR_NO_ERR(218)
                                                                                                                                ISR_NO_ERR(219) ISR_NO_ERR(220) ISR_NO_ERR(221) ISR_NO_ERR(222) ISR_NO_ERR(223) ISR_NO_ERR(224) ISR_NO_ERR(225) ISR_NO_ERR(226) ISR_NO_ERR(227) ISR_NO_ERR(
                                                                                                                                    228) ISR_NO_ERR(229)
                                                                                                                                    ISR_NO_ERR(230) ISR_NO_ERR(231) ISR_NO_ERR(232) ISR_NO_ERR(233) ISR_NO_ERR(234) ISR_NO_ERR(235) ISR_NO_ERR(
                                                                                                                                        236) ISR_NO_ERR(237)
                                                                                                                                        ISR_NO_ERR(
                                                                                                                                            238) ISR_NO_ERR(239)
                                                                                                                                            ISR_NO_ERR(240) ISR_NO_ERR(241) ISR_NO_ERR(242) ISR_NO_ERR(243) ISR_NO_ERR(244) ISR_NO_ERR(245) ISR_NO_ERR(
                                                                                                                                                246) ISR_NO_ERR(247)
                                                                                                                                                ISR_NO_ERR(248) ISR_NO_ERR(
                                                                                                                                                    249) ISR_NO_ERR(250)
                                                                                                                                                    ISR_NO_ERR(251) ISR_NO_ERR(
                                                                                                                                                        252)
                                                                                                                                                        ISR_NO_ERR(
                                                                                                                                                            253)
                                                                                                                                                            ISR_NO_ERR(
                                                                                                                                                                254)
                                                                                                                                                                ISR_NO_ERR(
                                                                                                                                                                    255)

    // Array of all ISR handlers for easy initialization
    void *isr_table[256] = {
        isr0,   isr1,   isr2,   isr3,   isr4,   isr5,   isr6,   isr7,   isr8,
        isr9,   isr10,  isr11,  isr12,  isr13,  isr14,  isr15,  isr16,  isr17,
        isr18,  isr19,  isr20,  isr21,  isr22,  isr23,  isr24,  isr25,  isr26,
        isr27,  isr28,  isr29,  isr30,  isr31,  isr32,  isr33,  isr34,  isr35,
        isr36,  isr37,  isr38,  isr39,  isr40,  isr41,  isr42,  isr43,  isr44,
        isr45,  isr46,  isr47,  isr48,  isr49,  isr50,  isr51,  isr52,  isr53,
        isr54,  isr55,  isr56,  isr57,  isr58,  isr59,  isr60,  isr61,  isr62,
        isr63,  isr64,  isr65,  isr66,  isr67,  isr68,  isr69,  isr70,  isr71,
        isr72,  isr73,  isr74,  isr75,  isr76,  isr77,  isr78,  isr79,  isr80,
        isr81,  isr82,  isr83,  isr84,  isr85,  isr86,  isr87,  isr88,  isr89,
        isr90,  isr91,  isr92,  isr93,  isr94,  isr95,  isr96,  isr97,  isr98,
        isr99,  isr100, isr101, isr102, isr103, isr104, isr105, isr106, isr107,
        isr108, isr109, isr110, isr111, isr112, isr113, isr114, isr115, isr116,
        isr117, isr118, isr119, isr120, isr121, isr122, isr123, isr124, isr125,
        isr126, isr127, isr128, isr129, isr130, isr131, isr132, isr133, isr134,
        isr135, isr136, isr137, isr138, isr139, isr140, isr141, isr142, isr143,
        isr144, isr145, isr146, isr147, isr148, isr149, isr150, isr151, isr152,
        isr153, isr154, isr155, isr156, isr157, isr158, isr159, isr160, isr161,
        isr162, isr163, isr164, isr165, isr166, isr167, isr168, isr169, isr170,
        isr171, isr172, isr173, isr174, isr175, isr176, isr177, isr178, isr179,
        isr180, isr181, isr182, isr183, isr184, isr185, isr186, isr187, isr188,
        isr189, isr190, isr191, isr192, isr193, isr194, isr195, isr196, isr197,
        isr198, isr199, isr200, isr201, isr202, isr203, isr204, isr205, isr206,
        isr207, isr208, isr209, isr210, isr211, isr212, isr213, isr214, isr215,
        isr216, isr217, isr218, isr219, isr220, isr221, isr222, isr223, isr224,
        isr225, isr226, isr227, isr228, isr229, isr230, isr231, isr232, isr233,
        isr234, isr235, isr236, isr237, isr238, isr239, isr240, isr241, isr242,
        isr243, isr244, isr245, isr246, isr247, isr248, isr249, isr250, isr251,
        isr252, isr253, isr254, isr255,
};

// Initialize all IDT entries
void init_idt(void) {
  for (int i = 0; i < 256; i++) {
    set_idt(i, isr_table[i], 0);
  }
  load_idt(idt);
}
