#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 <uacpi/acpi.h>
#include <uacpi/status.h>
#include <uacpi/tables.h>
#include <uacpi/uacpi.h>

/**
 * @file
 *
 * @brief 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 nanoseconds for conversion
static uint64_t hpet_clock_nano;
/// Lock, which protects concurrent access. See \ref amd64/smp.c
static spin_lock_t hpet_lock = SPIN_LOCK_INIT;

/** @cond DOXYGEN_IGNORE */
extern void amd64_spin (void);
/** @endcond */

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

/// Write a HPET register. Assumes caller holds \ref hpet_lock
static void amd64_hpet_write (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;
  if (hpet_32bits)
    *(volatile uint32_t*)(hpet_vaddr + reg) = (value & 0xFFFFFFFF);
  else
    *(volatile uint64_t*)(hpet_vaddr + reg) = value;
}

/// Read current value of \ref HPET_MCVR register.
static uint64_t amd64_hpet_timestamp (void) { return amd64_hpet_read (HPET_MCVR); }

/** 
 * @brief Get current HPET timestamp in nanoseconds
 *
 * @param lock
 *  if true, hold \ref hpet_lock
 */
uint64_t amd64_hpet_current_nano (bool lock) {
  if (lock)
    spin_lock (&hpet_lock);

  uint64_t t = amd64_hpet_timestamp () * hpet_clock_nano;

  if (lock)
    spin_unlock (&hpet_lock);

  return t;
}

/// 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) {
  spin_lock (&hpet_lock);

  uint64_t start = amd64_hpet_timestamp ();
  uint64_t conv = us * 1000;
  while (((amd64_hpet_timestamp () - start) * hpet_clock_nano) < conv)
    __asm__ volatile ("pause" ::: "memory");

  spin_unlock (&hpet_lock);
}

/// 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");
    amd64_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 | MM_PD_RELOAD);

  hpet_32bits = (amd64_hpet_read (HPET_GCIDR) & (1 << 13)) ? 0 : 1;

  /* reset */
  amd64_hpet_write (HPET_GCR, 0);
  amd64_hpet_write (HPET_MCVR, 0);
  amd64_hpet_write (HPET_GCR, 1);

  uint64_t gcidr = amd64_hpet_read (HPET_GCIDR);
  if (hpet_32bits) {
    uint32_t low = (uint32_t)gcidr;
    uint32_t high = (uint32_t)amd64_hpet_read (HPET_GCIDR + 4);
    gcidr = (((uint64_t)high << 32) | low);
  }

  uint64_t period_fs = (gcidr >> 32);

  hpet_clock_nano = period_fs / 1000000;
}