#include <amd64/apic.h>
#include <amd64/gdt.h>
#include <amd64/intr.h>
#include <amd64/intr_defs.h>
#include <amd64/io.h>
#include <aux/compiler.h>
#include <irq/irq.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/liballoc.h>
#include <proc/reschedule.h>
#include <sys/debug.h>
#include <sys/smp.h>
#include <sys/spin.h>
#include <syscall/syscall.h>
#include <syscall_defs.h>

/* 8259 PIC defs. */
#define PIC1      0x20
#define PIC2      0xA0
#define PIC1_CMD  PIC1
#define PIC1_DATA (PIC1 + 1)
#define PIC2_CMD  PIC2
#define PIC2_DATA (PIC2 + 1)
#define PIC_EOI   0x20

#define ICW1_ICW4   0x01
#define ICW1_SINGLE 0x02
#define ICW1_INTVL4 0x04
#define ICW1_LEVEL  0x08
#define ICW1_INIT   0x10

#define ICW4_8086     0x01
#define ICW4_AUTO     0x02
#define ICW4_BUFSLAVE 0x08
#define ICW4_BUFMASER 0x0C
#define ICW4_SFNM     0x10

#define CASCADE_IRQ 2

/* IDT defs. */

#define IDT_ENTRIES_MAX 256

/* 64-bit <IDT entry structure: https://wiki.osdev.org/Interrupt_Descriptor_Table */
struct idt_entry {
  uint16_t intrlow;
  uint16_t kernel_cs;
  uint8_t ist;
  uint8_t attrs;
  uint16_t intrmid;
  uint32_t intrhigh;
  uint32_t resv;
} PACKED;

struct idt {
  uint16_t limit;
  uint64_t base;
} PACKED;

ALIGNED (16) static volatile struct idt_entry idt_entries[IDT_ENTRIES_MAX];
static volatile struct idt idt;

/* Remaps and disables old 8259 PIC, since we'll be using APIC. */
static void pic_init (void) {
  outb (PIC1_CMD, (ICW1_INIT | ICW1_ICW4));
  io_wait ();
  outb (PIC2_CMD, (ICW1_INIT | ICW1_ICW4));
  io_wait ();

  outb (PIC1_DATA, 0x20);
  io_wait ();
  outb (PIC2_DATA, 0x28);
  io_wait ();

  outb (PIC1_DATA, (1 << CASCADE_IRQ));
  io_wait ();
  outb (PIC2_DATA, 2);
  io_wait ();

  outb (PIC1_DATA, ICW4_8086);
  io_wait ();
  outb (PIC2_DATA, ICW4_8086);
  io_wait ();

  /* Disable */
  outb (PIC1_DATA, 0xFF);
  io_wait ();
  outb (PIC2_DATA, 0xFF);
  io_wait ();
}

/* Set IDT entry */
static void idt_set (volatile struct idt_entry* ent, uint64_t handler, uint8_t flags, uint8_t ist) {
  ent->intrlow = (handler & 0xFFFF);
  ent->kernel_cs = GDT_KCODE;
  ent->ist = ist;
  ent->attrs = flags;
  ent->intrmid = ((handler >> 16) & 0xFFFF);
  ent->intrhigh = ((handler >> 32) & 0xFFFFFFFF);
  ent->resv = 0;
}

/* Load the IDT */
void idt_load (void) { __asm__ volatile ("lidt %0" ::"m"(idt)); }

/* Initialize IDT entries */
static void idt_init (void) {
  memset ((void*)idt_entries, 0, sizeof (idt_entries));

#define IDT_ENTRY(n, ist)                                                                          \
  extern void intr##n (void);                                                                      \
  idt_set (&idt_entries[(n)], (uint64_t)&intr##n, 0x8E, (ist))
  /* clang-format off */
  IDT_ENTRY (0, 0); IDT_ENTRY (1, 0); IDT_ENTRY (2, 0); IDT_ENTRY (3, 0);
  IDT_ENTRY (4, 0); IDT_ENTRY (5, 0); IDT_ENTRY (6, 0); IDT_ENTRY (7, 0);
  IDT_ENTRY (8, 0); IDT_ENTRY (9, 0); IDT_ENTRY (10, 0); IDT_ENTRY (11, 0);
  IDT_ENTRY (12, 0); IDT_ENTRY (13, 0); IDT_ENTRY (14, 0); IDT_ENTRY (15, 0);
  IDT_ENTRY (16, 0); IDT_ENTRY (17, 0); IDT_ENTRY (18, 0); IDT_ENTRY (19, 0);
  IDT_ENTRY (20, 0); IDT_ENTRY (21, 0); IDT_ENTRY (22, 0); IDT_ENTRY (23, 0);
  IDT_ENTRY (24, 0); IDT_ENTRY (25, 0); IDT_ENTRY (26, 0); IDT_ENTRY (27, 0);
  IDT_ENTRY (28, 0); IDT_ENTRY (29, 0); IDT_ENTRY (30, 0); IDT_ENTRY (31, 0);
  IDT_ENTRY (32, 1); IDT_ENTRY (33, 1); IDT_ENTRY (34, 1); IDT_ENTRY (35, 1);
  IDT_ENTRY (36, 1); IDT_ENTRY (37, 1); IDT_ENTRY (38, 1); IDT_ENTRY (39, 1);
  IDT_ENTRY (40, 1); IDT_ENTRY (41, 1); IDT_ENTRY (42, 1); IDT_ENTRY (43, 1);
  IDT_ENTRY (44, 1); IDT_ENTRY (45, 1); IDT_ENTRY (46, 1); IDT_ENTRY (47, 1);

  IDT_ENTRY (SCHED_PREEMPT_TIMER, 1);
  IDT_ENTRY (CPU_REQUEST_SCHED, 1);
  IDT_ENTRY (CPU_SPURIOUS, 1);
  /* clang-format on */
#undef IDT_ENTRY

  idt.limit = sizeof (idt_entries) - 1;
  idt.base = (uint64_t)idt_entries;

  idt_load ();
}

/* Handle CPU exception and dump registers. If incoming CS has CPL3, kill the process. */
static void intr_exception (struct saved_regs* regs) {
  DEBUG ("cpu exception %lu (%lu)\n", regs->trap, regs->error);

  uint64_t cr2;
  __asm__ volatile ("movq %%cr2, %0" : "=r"(cr2));
  uint64_t cr3;
  __asm__ volatile ("movq %%cr3, %0" : "=r"(cr3));

  debugprintf ("r15=%016lx r14=%016lx r13=%016lx\n"
               "r12=%016lx r11=%016lx r10=%016lx\n"
               "r9 =%016lx r8 =%016lx rbp=%016lx\n"
               "rdi=%016lx rsi=%016lx rdx=%016lx\n"
               "rcx=%016lx rax=%016lx trp=%016lx\n"
               "err=%016lx rip=%016lx cs =%016lx\n"
               "rfl=%016lx rsp=%016lx ss =%016lx\n"
               "cr2=%016lx cr3=%016lx rbx=%016lx\n",
               regs->r15, regs->r14, regs->r13, regs->r12, regs->r11, regs->r10, regs->r9, regs->r8,
               regs->rbp, regs->rdi, regs->rsi, regs->rdx, regs->rcx, regs->rax, regs->trap,
               regs->error, regs->rip, regs->cs, regs->rflags, regs->rsp, regs->ss, cr2, cr3,
               regs->rbx);

  if (regs->cs == (GDT_UCODE | 0x03)) {
    struct reschedule_ctx rctx = {.reschedule = false, .cpu = NULL};
    proc_kill (thiscpu->proc_current, &rctx);

    if (rctx.reschedule)
      cpu_request_sched (rctx.cpu);
  } else {
    __asm__ volatile ("cli");
    spin ();
  }
}

/* Handle incoming interrupt, dispatch IRQ handlers. */
void intr_handler (void* stack_ptr) {
  load_kernel_cr3 ();

  struct saved_regs* regs = stack_ptr;

  if (regs->trap <= 31) {
    intr_exception (regs);
  } else {
    spin_lock (&thiscpu->lock);
    struct proc* proc_current = thiscpu->proc_current;
    spin_lock (&proc_current->lock);

    memcpy (&proc_current->pdata.regs, regs, sizeof (struct saved_regs));

    spin_unlock (&proc_current->lock);
    spin_unlock (&thiscpu->lock);

    lapic_eoi ();

    struct irq* irq = irq_find (regs->trap);

    if (irq == NULL)
      return;

    struct reschedule_ctx rctx = {.reschedule = false, .cpu = NULL};
    irq->func (irq->arg, stack_ptr, &rctx);

    if (rctx.reschedule)
      cpu_request_sched (rctx.cpu);
  }
}

/* Initialize interrupts */
void intr_init (void) {
  pic_init ();
  idt_init ();
}