#include <amd64/hpet.h>
#include <libk/std.h>
#include <limine/requests.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/mm.h>
#include <sys/spin.h>
#include <uacpi/acpi.h>
#include <uacpi/status.h>
#include <uacpi/tables.h>
#include <uacpi/uacpi.h>

/*
 * HPET (High Precision Event Timer) driver code. See more at https://wiki.osdev.org/HPET
 */

/* HPET Main Counter Value Register */
#define HPET_MCVR 0xF0
/* HPET General Configuration Register */
#define HPET_GCR 0x10
/* HPET General Capabilities and ID Register */
#define HPET_GCIDR 0x00

/* Set whether we sould use 32-bit or 64-bit reads/writes */
static bool hpet_32bits = 1;
/* Physical address for HPET MMIO */
static uintptr_t hpet_paddr;
/* HPET period in femtoseconds */
static uint64_t hpet_period_fs;
/* Lock, which protects concurrent access. See amd64/smp.c */
static spin_lock_t hpet_lock = SPIN_LOCK_INIT;

/* Read a HPET register. Assumes caller holds hpet_lock */
static uint64_t amd64_hpet_read64 (uint32_t reg) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t hpet_vaddr = hpet_paddr + (uintptr_t)hhdm->offset;
  return *(volatile uint64_t*)(hpet_vaddr + reg);
}

static uint32_t amd64_hpet_read32 (uint32_t reg) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t hpet_vaddr = hpet_paddr + (uintptr_t)hhdm->offset;
  return *(volatile uint32_t*)(hpet_vaddr + reg);
}

/* Write a HPET register. Assumes caller holds hpet_lock */
static void amd64_hpet_write64 (uint32_t reg, uint64_t value) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t hpet_vaddr = hpet_paddr + (uintptr_t)hhdm->offset;
  *(volatile uint64_t*)(hpet_vaddr + reg) = value;
}

static void amd64_hpet_write32 (uint32_t reg, uint32_t value) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  uintptr_t hpet_vaddr = hpet_paddr + (uintptr_t)hhdm->offset;
  *(volatile uint32_t*)(hpet_vaddr + reg) = value;
}

/* Read current value of HPET_MCVR register. */

static uint64_t amd64_hpet_read_counter (void) {
  uint64_t value;
  spin_lock_ctx_t ctxhrc;

  spin_lock (&hpet_lock, &ctxhrc);

  if (!hpet_32bits)
    value = amd64_hpet_read64 (HPET_MCVR);
  else {
    uint32_t hi1, lo, hi2;
    do {
      hi1 = amd64_hpet_read32 (HPET_MCVR + 4);
      lo = amd64_hpet_read32 (HPET_MCVR + 0);
      hi2 = amd64_hpet_read32 (HPET_MCVR + 4);
    } while (hi1 != hi2);

    value = ((uint64_t)hi1 << 32) | lo;
  }

  spin_unlock (&hpet_lock, &ctxhrc);

  return value;
}

static void amd64_hpet_write_counter (uint64_t value) {
  spin_lock_ctx_t ctxhwc;

  spin_lock (&hpet_lock, &ctxhwc);

  if (!hpet_32bits)
    amd64_hpet_write64 (HPET_MCVR, value);
  else {
    amd64_hpet_write32 (HPET_MCVR, (uint32_t)value);
    amd64_hpet_write32 (HPET_MCVR + 4, (uint32_t)(value >> 32));
  }

  spin_unlock (&hpet_lock, &ctxhwc);
}

/* Sleep for a given amount of microseconds. This time can last longer due to \ref hpet_lock being
 * held. */
void amd64_hpet_sleep_micro (uint64_t us) {
  if (hpet_period_fs == 0)
    return;

  uint64_t ticks_to_wait = (us * 1000ULL) / (hpet_period_fs / 1000000ULL);
  uint64_t start = amd64_hpet_read_counter ();

  for (;;) {
    uint64_t now = amd64_hpet_read_counter ();

    if ((now - start) >= ticks_to_wait)
      break;

    __asm__ volatile ("pause" ::: "memory");
  }
}

/* Initialize HPET */
void amd64_hpet_init (void) {
  struct uacpi_table hpet_table;
  uacpi_status status = uacpi_table_find_by_signature (ACPI_HPET_SIGNATURE, &hpet_table);
  if (status != UACPI_STATUS_OK) {
    DEBUG ("Could not find HPET table!\n");
    spin ();
  }

  struct acpi_hpet* hpet = (struct acpi_hpet*)hpet_table.virt_addr;
  hpet_paddr = (uintptr_t)hpet->address.address;

  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  mm_map_kernel_page (hpet_paddr, (uintptr_t)hhdm->offset + hpet_paddr, MM_PG_PRESENT | MM_PG_RW);

  uint64_t caps = amd64_hpet_read64 (HPET_GCIDR);
  hpet_32bits = (caps & (1 << 13)) ? 0 : 1;

  hpet_period_fs = (uint32_t)(caps >> 32);

  amd64_hpet_write64 (HPET_GCR, 0);
  amd64_hpet_write_counter (0);
  amd64_hpet_write64 (HPET_GCR, 1);
}