#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <amd64/msr-index.h>
#include <amd64/msr.h>
#include <libk/std.h>
#include <limine/requests.h>
#include <sys/debug.h>
#include <sys/mm.h>
#include <sys/spin.h>
#include <sys/time.h>
#include <uacpi/acpi.h>
#include <uacpi/status.h>
#include <uacpi/tables.h>
#include <uacpi/uacpi.h>

#define IOAPICS_MAX                 24
#define INTERRUPT_SRC_OVERRIDES_MAX 24

/* ID of Local APIC */
#define LAPIC_ID 0x20
/* End of interrupt register */
#define LAPIC_EOI 0xB0
/* Spurious interrupt vector register */
#define LAPIC_SIVR 0xF0
/* Interrupt command register */
#define LAPIC_ICR 0x300
/* LVT timer register */
#define LAPIC_LVTTR 0x320
/* Timer initial count register */
#define LAPIC_TIMICT 0x380
/* Timer current count register */
#define LAPIC_TIMCCT 0x390
/* Divide config register */
#define LAPIC_DCR 0x3E0

/* Table of IOAPICS */
static struct acpi_madt_ioapic apics[IOAPICS_MAX];
/* Table of interrupt source overrides */
/* clang-format off */
static struct acpi_madt_interrupt_source_override intr_src_overrides[INTERRUPT_SRC_OVERRIDES_MAX];
/* clang-format on */
/* Count of actual IOAPIC entries */
static size_t ioapic_entries = 0;
/* Count of actual interrupt source overrides */
static size_t intr_src_override_entries = 0;
/* Local APIC MMIO base address. It comes from MSR_APIC_BASE */
static uintptr_t lapic_mmio_base = 0;

/* Read IOAPIC */
static uint32_t amd64_ioapic_read (uintptr_t vaddr, uint32_t reg) {
  *(volatile uint32_t*)vaddr = reg;
  return *(volatile uint32_t*)(vaddr + 0x10);
}

/* Write IOAPIC */
static void amd64_ioapic_write (uintptr_t vaddr, uint32_t reg, uint32_t value) {
  *(volatile uint32_t*)vaddr = reg;
  *(volatile uint32_t*)(vaddr + 0x10) = value;
}

/* Find an IOAPIC corresposting to provided IRQ */
static struct acpi_madt_ioapic* amd64_ioapic_find (uint8_t irq) {
  struct acpi_madt_ioapic* apic = NULL;
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  for (size_t i = 0; i < ioapic_entries; i++) {
    apic = &apics[i];
    uint32_t version = amd64_ioapic_read ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, 1);
    uint32_t max = ((version >> 16) & 0xFF);

    if ((irq >= apic->gsi_base) && (irq <= (apic->gsi_base + max)))
      return apic;
  }

  return NULL;
}

/*
 * Route IRQ to an IDT entry of a given Local APIC.
 *
 * vec - Interrupt vector number, which will be delivered to the CPU.
 * irq -Legacy IRQ number to be routed. Can be changed by an interrupt source override
 *  into a different GSI.
 * flags - IOAPIC redirection flags.
 * lapic_id - Local APIC that will receive the interrupt.
 */
void amd64_ioapic_route_irq (uint8_t vec, uint8_t irq, uint64_t flags, uint64_t lapic_id) {
  struct acpi_madt_ioapic* apic = NULL;
  struct acpi_madt_interrupt_source_override* override;
  bool found_override = false;
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  for (size_t i = 0; i < intr_src_override_entries; i++) {
    override = &intr_src_overrides[i];
    if (override->source == irq) {
      found_override = true;
      break;
    }
  }

  uint64_t calc_flags = (lapic_id << 56) | (flags) | (vec & 0xFF);

  if (found_override) {
    uint8_t polarity = ((override->flags & 0x03) == 0x03) ? 1 : 0;
    uint8_t mode = (((override->flags >> 2) & 0x03) == 0x03) ? 1 : 0;
    calc_flags |= (uint64_t)mode << 15;
    calc_flags |= (uint64_t)polarity << 13;

    calc_flags |= flags;
  } else {
    calc_flags |= flags;
  }

  apic = amd64_ioapic_find (irq);

  if (apic == NULL)
    return;

  uint32_t irq_reg = ((irq - apic->gsi_base) * 2) + 0x10;

  amd64_ioapic_write ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg,
                      (uint32_t)calc_flags);

  amd64_ioapic_write ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg + 1,
                      (uint32_t)(calc_flags >> 32));
}

/* Mask a given IRQ */
void amd64_ioapic_mask (uint8_t irq) {
  struct acpi_madt_ioapic* apic;
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  apic = amd64_ioapic_find (irq);

  if (apic == NULL)
    return;

  uint32_t irq_reg = ((irq - apic->gsi_base) * 2) + 0x10;

  uint32_t value = amd64_ioapic_read ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg);
  amd64_ioapic_write ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg,
                      value | (1 << 16));
}

/* Unmask a given IRQ */
void amd64_ioapic_unmask (uint8_t irq) {
  struct acpi_madt_ioapic* apic;
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  apic = amd64_ioapic_find (irq);

  if (apic == NULL)
    return;

  uint32_t irq_reg = ((irq - apic->gsi_base) * 2) + 0x10;

  uint32_t value = amd64_ioapic_read ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg);
  amd64_ioapic_write ((uintptr_t)hhdm->offset + (uintptr_t)apic->address, irq_reg,
                      value & ~(1 << 16));
}

/* Find and initialize the IOAPIC */
void amd64_ioapic_init (void) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct uacpi_table apic_table;
  uacpi_status status = uacpi_table_find_by_signature (ACPI_MADT_SIGNATURE, &apic_table);
  if (status != UACPI_STATUS_OK) {
    DEBUG ("Could not find MADT table!\n");
    spin ();
  }

  struct acpi_madt* apic = (struct acpi_madt*)apic_table.virt_addr;
  struct acpi_entry_hdr* current = (struct acpi_entry_hdr*)apic->entries;

  for (;;) {
    if ((uintptr_t)current >=
        ((uintptr_t)apic->entries + apic->hdr.length - sizeof (struct acpi_madt)))
      break;

    switch (current->type) {
    case ACPI_MADT_ENTRY_TYPE_IOAPIC: {
      struct acpi_madt_ioapic* ioapic = (struct acpi_madt_ioapic*)current;
      mm_map_kernel_page ((uintptr_t)ioapic->address,
                          (uintptr_t)hhdm->offset + (uintptr_t)ioapic->address,
                          MM_PG_PRESENT | MM_PG_RW | MM_PD_RELOAD);
      apics[ioapic_entries++] = *ioapic;
    } break;
    case ACPI_MADT_ENTRY_TYPE_INTERRUPT_SOURCE_OVERRIDE: {
      struct acpi_madt_interrupt_source_override* override =
          (struct acpi_madt_interrupt_source_override*)current;
      intr_src_overrides[intr_src_override_entries++] = *override;
    } break;
    }

    current = (struct acpi_entry_hdr*)((uintptr_t)current + current->length);
  }
}

/* Get MMIO base of Local APIC */
static uintptr_t amd64_lapic_base (void) { return lapic_mmio_base; }

/* Write Local APIC */
static void amd64_lapic_write (uint32_t reg, uint32_t value) {
  *(volatile uint32_t*)(amd64_lapic_base () + reg) = value;
}

/* Read Local APIC */
static uint32_t amd64_lapic_read (uint32_t reg) {
  return *(volatile uint32_t*)(amd64_lapic_base () + reg);
}

/* Get ID of Local APIC */
uint32_t amd64_lapic_id (void) { return amd64_lapic_read (LAPIC_ID) >> 24; }

/* Send End of interrupt command to Local APIC */
void amd64_lapic_eoi (void) { amd64_lapic_write (LAPIC_EOI, 0); }

/* Set initial counter value in Local APIC timer */
void amd64_lapic_tick (uint32_t tick) { amd64_lapic_write (LAPIC_TIMICT, tick); }

/*
 * Calibrate Local APIC to send interrupts in a set interval.
 *
 * us - Period length in microseconds
 */
static uint32_t amd64_lapic_calibrate (uint32_t us) {
  amd64_lapic_write (LAPIC_DCR, 0x0B);

  amd64_lapic_write (LAPIC_LVTTR, SCHED_PREEMPT_TIMER | (1 << 16));

  amd64_lapic_write (LAPIC_TIMICT, 0xFFFFFFFF);

  sleep_micro (us);

  uint32_t ticks = 0xFFFFFFFF - amd64_lapic_read (LAPIC_TIMCCT);

  return ticks;
}

/*
 * Starts a Local APIC, configures LVT timer to send interrupts at SCHED_PREEMPT_TIMER.
 *
 * ticks - Initial tick count
 */
static void amd64_lapic_start (uint32_t ticks) {
  amd64_lapic_write (LAPIC_DCR, 0x0B);

  amd64_lapic_write (LAPIC_LVTTR, SCHED_PREEMPT_TIMER | (1 << 17));

  amd64_lapic_write (LAPIC_TIMICT, ticks);
}

/*
 * Initialize Local APIC, configure to send timer interrupts at a given period. See
 * amd64_lapic_calibrate and amd64_lapic_start.
 */
uint64_t amd64_lapic_init (uint32_t us) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  amd64_wrmsr (MSR_APIC_BASE, amd64_rdmsr (MSR_APIC_BASE) | (1 << 11));

  uintptr_t lapic_paddr = amd64_rdmsr (MSR_APIC_BASE) & 0xFFFFF000;
  lapic_mmio_base = lapic_paddr + (uintptr_t)hhdm->offset;

  mm_map_kernel_page (lapic_paddr, lapic_mmio_base,
                      MM_PG_PRESENT | MM_PG_RW | MM_PD_LOCK | MM_PD_RELOAD);

  amd64_lapic_write (LAPIC_SIVR, 0xFF | (1 << 8));

  uint32_t ticks = amd64_lapic_calibrate (us);

  amd64_lapic_start (ticks);

  return ticks;
}

/*
 * Send an IPI to a given Local APIC. This till invoke an IDT stub located at vec.
 *
 * lapic_id - Target Local APIC
 * vec - Interrupt vector/IDT stub, which will be invoked by the IPI.
 */
void amd64_lapic_ipi (uint8_t lapic_id, uint8_t vec) {
  amd64_lapic_write (LAPIC_ICR + 0x10, (lapic_id << 24));
  amd64_lapic_write (LAPIC_ICR, vec);
}
