#include "common.h"
#include "map.h"
#include "print.h"

// stub of a PIC layout, we don't need the whole thing
// as we only care about initializing it to disable it
enum {
  PIC_COMM_MASTER = 0x20,
  PIC_COMM_SLAVE = 0xA0,
  PIC_DATA_MASTER = 0x21,
  PIC_DATA_SLAVE = 0xA1,
};

#define IOAPIC_REG_REDIR(n) (0x10 + 2 * (n))

enum {
  IOAPIC_REG_SELECT = 0x00,
  IOAPIC_REG_WINDOW = 0x10,
};

// I/O APIC redirection entry bits
#define IOAPIC_DST_PHYS (0 << 11)
#define IOAPIC_DST_LOGICAL (1 << 11)
#define IOAPIC_MASKED (1 << 16)
#define IOAPIC_LVL_TRIGG (1 << 15)
#define IOAPIC_EDGE_TRIGG (0 << 15)
#define IOAPIC_ACTIVE_LOW (1 << 13)
#define IOAPIC_ACTIVE_HIGH (0 << 13)

static void pic_init() {
  // Start initialization sequence
  outb(PIC_COMM_MASTER, 0x11); // ICW1: Initialize in cascading mode
  outb(PIC_DATA_MASTER, 0x20); // ICW2: Master PIC vector offset
  outb(PIC_DATA_MASTER, 0x04); // ICW3: Slave PIC at IRQ2
  outb(PIC_DATA_MASTER, 0x01); // ICW4: 8086/88 (MCS-80 compatible)

  outb(PIC_COMM_SLAVE, 0x11); // ICW1: Initialize slave PIC
  outb(PIC_DATA_SLAVE, 0x28); // ICW2: Slave PIC vector offset
  outb(PIC_DATA_SLAVE, 0x02); // ICW3: It's connected to master PIC
  outb(PIC_DATA_SLAVE, 0x01); // ICW4: 8086/88 mode

  // Mask all interrupts
  outb(PIC_DATA_MASTER, 0xFF);
  outb(PIC_DATA_SLAVE, 0xFF);
}

static void pic_disable() {
  pic_init();
  outb(PIC_DATA_MASTER, ~0);
  outb(PIC_DATA_SLAVE, ~0);
}

static inline u64 rdmsr(u32 msr) {
  u32 hi, lo;
  asm volatile("rdmsr" : "=a"(lo), "=d"(hi) : "c"(msr));

  return (((u64)hi) << 32) | (u64)lo;
}

static inline void wrmsr(u32 msr, u64 val) {
  u32 eax = (u32)val;
  u32 edx = (u32)(val >> 32);
  asm volatile("wrmsr" : : "c"(msr), "a"(eax), "d"(edx));
}

static u64 apic_base;

u64 apic_get_local_base() {
  const u8 APIC_BASE_MSR = 0x1B;

  u64 msr = rdmsr(APIC_BASE_MSR);

  return (apic_base = (msr & 0xfffff000));
}

static u64 apic_set_local_base(u64 base) {
  const u8 APIC_BASE_MSR = 0x1B;

  u64 msr = rdmsr(APIC_BASE_MSR);
  msr &= 0xFFF; // keep flags
  msr |= (base & 0xFFFFF000);
  msr |= (1 << 11); // to ensure APIC is enabled

  wrmsr(APIC_BASE_MSR, msr);

  return base & 0xfffff000;
}

u32 *apic_reg(u32 reg) { return (u32 *)(hhdm_offset + apic_base + reg); }

void send_eoi() {
  volatile u32 *reg = apic_reg(0xB0);
  *reg = 0;
}

void apic_enable() {
  pic_disable();
  print("\nPIC Disabled\n");

  apic_set_local_base(apic_get_local_base());

  if (!map_mmio((u64)apic_base, 1 << 12)) {
    print("Failed to map APIC registers!\n");
    hang();
  }

  volatile u32 *spurious = apic_reg(0xF0);
  // set enable flag
  *spurious |= 0x100;

  // TPR = 0 (accept all interrupts)
  volatile u32 *tpr = apic_reg(0x80);
  *tpr = 0;

  // Mask all LVT antries except LINT0 (0x21)

  volatile u32 *timer = apic_reg(0x320);
  *timer = (1 << 16);
  volatile u32 *LINT0 = apic_reg(0x350);
  *LINT0 = (7 << 8);
  //  *LINT0 = (1 << 16);
  volatile u32 *LINT1 = apic_reg(0x360);
  *LINT1 = (1 << 16);
  volatile u32 *Error = apic_reg(0x370);
  *Error = (1 << 16);
}

static volatile u32 *ioapic_base;

void ioapic_write(u8 reg, u32 value) {
  ioapic_base[IOAPIC_REG_SELECT / 4] = reg;
  ioapic_base[IOAPIC_REG_WINDOW / 4] = value;
}

u32 ioapic_read(u8 reg) {
  ioapic_base[IOAPIC_REG_SELECT / 4] = reg;
  return ioapic_base[IOAPIC_REG_WINDOW / 4];
}

void ioapic_init() {
  // simply assume it is at the default address
  ioapic_base = map_mmio(0xFEC00000, 4096);
  print("IOAPIC BASE: ");
  print64((u64)ioapic_base);
  print("\nIOAPIC Version:");
  ioapic_base[0] = 0x01;
  print32(ioapic_base[4]);
  print("\n");

  // mask all IRQs
  for (size_t i = 0; i < 24; i++) {
    u8 vect = 0x20 + i;
    u32 low = vect | IOAPIC_MASKED | IOAPIC_DST_PHYS | IOAPIC_EDGE_TRIGG |
              IOAPIC_ACTIVE_HIGH;
    u32 high = 0;

    ioapic_write(0x10 + i * 2, low);
    ioapic_write(0x10 + i * 2 + 1, high);
  }

  // setup keyboard by selecting IRQ1
  u8 irq = 1;
  u8 vect = 0x21;
  u8 cpu = 0;

  u32 low = ioapic_read(IOAPIC_REG_REDIR(irq));
  u32 high = ioapic_read(IOAPIC_REG_REDIR(irq) + 1);

  // Configure keyboard - UNMASKED (no IOAPIC_MASKED bit!)
  low = vect | IOAPIC_DST_PHYS | IOAPIC_EDGE_TRIGG | IOAPIC_ACTIVE_HIGH;
  high = (cpu << 24);

  ioapic_write(IOAPIC_REG_REDIR(irq), low);
  ioapic_write(IOAPIC_REG_REDIR(irq) + 1, high);
}
