#include <amd64/apic.h>
#include <amd64/intr.h>
#include <amd64/io.h>
#include <aux/compiler.h>
#include <irq/irq.h>
#include <libk/std.h>
#include <libk/string.h>
#include <sys/debug.h>
#include <sys/irq.h>
#include <sys/smp.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

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;

extern void amd64_spin (void);

/* Remaps and disables old 8259 PIC, since we'll be using APIC. */
static void amd64_init_pic (void) {
#define IO_OP(fn, ...)                                                                             \
  fn (__VA_ARGS__);                                                                                \
  amd64_io_wait ()

  IO_OP (amd64_io_outb, PIC1_CMD, (ICW1_INIT | ICW1_ICW4));
  IO_OP (amd64_io_outb, PIC2_CMD, (ICW1_INIT | ICW1_ICW4));

  IO_OP (amd64_io_outb, PIC1_DATA, 0x20);
  IO_OP (amd64_io_outb, PIC2_DATA, 0x28);

  IO_OP (amd64_io_outb, PIC1_DATA, (1 << CASCADE_IRQ));
  IO_OP (amd64_io_outb, PIC2_DATA, 2);

  IO_OP (amd64_io_outb, PIC1_DATA, ICW4_8086);
  IO_OP (amd64_io_outb, PIC2_DATA, ICW4_8086);

  /* Disable */
  IO_OP (amd64_io_outb, PIC1_DATA, 0xFF);
  IO_OP (amd64_io_outb, PIC2_DATA, 0xFF);

#undef IO_OP
}

static void amd64_idt_set (volatile struct idt_entry* ent, uint64_t handler, uint8_t flags,
                           uint8_t ist) {
  ent->intrlow = (handler & 0xFFFF);
  ent->kernel_cs = 0x08; // GDT_KCODE (init.c)
  ent->ist = ist;
  ent->attrs = flags;
  ent->intrmid = ((handler >> 16) & 0xFFFF);
  ent->intrhigh = ((handler >> 32) & 0xFFFFFFFF);
  ent->resv = 0;
}

void amd64_load_idt (void) { __asm__ volatile ("lidt %0" ::"m"(idt)); }

static void amd64_idt_init (void) {
  memset ((void*)idt_entries, 0, sizeof (idt_entries));

#define IDT_ENTRY(n, ist)                                                                          \
  extern void amd64_intr##n (void);                                                                \
  amd64_idt_set (&idt_entries[(n)], (uint64_t)&amd64_intr##n, 0x8E, (ist))
  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);
#undef IDT_ENTRY

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

  amd64_load_idt ();
}

static void amd64_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);

  amd64_spin ();
}

void amd64_intr_handler (void* stack_ptr) {
  struct saved_regs* regs = stack_ptr;

  if (regs->trap <= 31) {
    amd64_intr_exception (regs);
  } else if (regs->trap >= 32) {
    amd64_lapic_eoi ();

    __asm__ volatile ("sti");

    irq_invoke_each (regs->trap);

    __asm__ volatile ("cli");
  } else {
    DEBUG ("unknown trap %lu\n", regs->trap);
  }
}

void amd64_intr_init (void) {
  amd64_init_pic ();
  amd64_idt_init ();
}

/* Aux. */

static uint64_t amd64_irq_save_flags (void) {
  uint64_t rflags;
  __asm__ volatile ("pushfq; cli; popq %0" : "=r"(rflags)::"memory", "cc");
  return rflags;
}

static void amd64_irq_restore_flags (uint64_t rflags) {
  if (rflags & (1ULL << 9))
    __asm__ volatile ("sti");
}

void irq_save (void) {
  /* before smp init. */
  if (thiscpu == NULL)
    return;

  int prev = atomic_fetch_add_explicit (&thiscpu->irq_ctx.nesting, 1, memory_order_acq_rel);
  if (prev == 0)
    thiscpu->irq_ctx.rflags = amd64_irq_save_flags ();
}

void irq_restore (void) {
  /* before smp init. */
  if (thiscpu == NULL)
    return;

  int prev = atomic_fetch_sub_explicit (&thiscpu->irq_ctx.nesting, 1, memory_order_acq_rel);
  if (prev == 1)
    amd64_irq_restore_flags (thiscpu->irq_ctx.rflags);
}