#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 <sync/rw_spin_lock.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

#define DIVIDER_VALUE 0x0B

struct ioapic {
  struct acpi_madt_ioapic table_data;
  rw_spin_lock_t lock;
  uintptr_t mmio_base;
};

/* Table of IOAPICS */
static struct ioapic ioapics[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;

static spin_lock_t lapic_calibration_lock = SPIN_LOCK_INIT;

/* Read IOAPIC */
static uint32_t amd64_ioapic_read (struct ioapic* ioapic, uint32_t reg) {
  spin_lock_ctx_t ctxioar;

  rw_spin_read_lock (&ioapic->lock, &ctxioar);
  *(volatile uint32_t*)ioapic->mmio_base = reg;
  uint32_t ret = *(volatile uint32_t*)(ioapic->mmio_base + 0x10);
  rw_spin_read_unlock (&ioapic->lock, &ctxioar);
  return ret;
}

/* Write IOAPIC */
static void amd64_ioapic_write (struct ioapic* ioapic, uint32_t reg, uint32_t value) {
  spin_lock_ctx_t ctxioaw;

  rw_spin_write_lock (&ioapic->lock, &ctxioaw);
  *(volatile uint32_t*)ioapic->mmio_base = reg;
  *(volatile uint32_t*)(ioapic->mmio_base + 0x10) = value;
  rw_spin_write_unlock (&ioapic->lock, &ctxioaw);
}

/* Find an IOAPIC corresposting to provided IRQ */
static struct ioapic* amd64_ioapic_find (uint32_t irq) {
  struct ioapic* ioapic = NULL;

  for (size_t i = 0; i < ioapic_entries; i++) {
    ioapic = &ioapics[i];
    uint32_t version = amd64_ioapic_read (ioapic, 1);
    uint32_t max = ((version >> 16) & 0xFF);

    if ((irq >= ioapic->table_data.gsi_base) && (irq <= (ioapic->table_data.gsi_base + max)))
      return ioapic;
  }

  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 (uint32_t vec, uint32_t irq, uint64_t flags, uint64_t lapic_id) {
  struct ioapic* ioapic = NULL;
  struct acpi_madt_interrupt_source_override* override;
  bool found_override = false;

  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) {
    uint32_t polarity = ((override->flags & 0x03) == 0x03) ? 1 : 0;
    uint32_t mode = (((override->flags >> 2) & 0x03) == 0x03) ? 1 : 0;
    calc_flags |= (uint64_t)mode << 15;
    calc_flags |= (uint64_t)polarity << 13;
  }

  uint32_t gsi = found_override ? override->gsi : irq;

  ioapic = amd64_ioapic_find (gsi);

  if (ioapic == NULL)
    return;

  uint32_t irq_reg = ((gsi - ioapic->table_data.gsi_base) * 2) + 0x10;

  amd64_ioapic_write (ioapic, irq_reg + 1, (uint32_t)(calc_flags >> 32));
  amd64_ioapic_write (ioapic, irq_reg, (uint32_t)calc_flags);
}

/* 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_table_data = (struct acpi_madt_ioapic*)current;
      mm_map_kernel_page ((uintptr_t)ioapic_table_data->address,
                          (uintptr_t)hhdm->offset + (uintptr_t)ioapic_table_data->address,
                          MM_PG_PRESENT | MM_PG_RW | MM_PD_RELOAD);
      ioapics[ioapic_entries++] = (struct ioapic){
          .lock = RW_SPIN_LOCK_INIT,
          .table_data = *ioapic_table_data,
          .mmio_base = ((uintptr_t)hhdm->offset + (uintptr_t)ioapic_table_data->address),
      };
    } 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 thiscpu->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); }

/*
 * Calibrate Local APIC to send interrupts in a set interval.
 *
 * us - Period length in microseconds
 */
static uint32_t amd64_lapic_calibrate (uint32_t us) {
  spin_lock_ctx_t ctxlacb;

  spin_lock (&lapic_calibration_lock, &ctxlacb);

  amd64_lapic_write (LAPIC_DCR, DIVIDER_VALUE);

  amd64_lapic_write (LAPIC_LVTTR, SCHED_PREEMPT_TIMER | (1 << 16));
  amd64_lapic_write (LAPIC_TIMICT, 0xFFFFFFFF);

  sleep_micro (us);

  amd64_lapic_write (LAPIC_LVTTR, SCHED_PREEMPT_TIMER | (0 << 16));
  uint32_t ticks = 0xFFFFFFFF - amd64_lapic_read (LAPIC_TIMCCT);
  DEBUG ("timer ticks = %u\n", ticks);

  spin_unlock (&lapic_calibration_lock, &ctxlacb);

  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, DIVIDER_VALUE);
  amd64_lapic_write (LAPIC_TIMICT, ticks);
  amd64_lapic_write (LAPIC_LVTTR, SCHED_PREEMPT_TIMER | (1 << 17));
}

/*
 * Initialize Local APIC, configure to send timer interrupts at a given period. See
 * amd64_lapic_calibrate and amd64_lapic_start.
 */
void 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;
  thiscpu->lapic_mmio_base = lapic_paddr + (uintptr_t)hhdm->offset;

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

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

  thiscpu->lapic_ticks = amd64_lapic_calibrate (us);
  amd64_lapic_start (thiscpu->lapic_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 (uint32_t lapic_id, uint32_t vec) {
  /* wait for previous IPI to finish */
  while (amd64_lapic_read (LAPIC_ICR) & (1 << 12)) {
    __asm__ volatile ("pause");
  }

  amd64_lapic_write (LAPIC_ICR + 0x10, (lapic_id << 24));
  amd64_lapic_write (LAPIC_ICR, vec | (1 << 14));
}