#include <zinnia/acpi.h>
#include <zinnia/assert.h>
#include <zinnia/clock.h>
#include <zinnia/cmdline.h>
#include <zinnia/initgraph.h>
#include <zinnia/mmio.h>
#include <zinnia/percpu.h>
#include <zinnia/pmap.h>
#include <zinnia/print.h>
#include <zinnia/utils.h>
#include <zinnia/vma.h>
#include <uacpi/acpi.h>
#include <uacpi/status.h>
#include <uacpi/tables.h>
#include <uacpi/uacpi.h>

static bool use_hpet = true;
static void hpet_option(bool use) {
    use_hpet = use;
}
CMDLINE_OPTION("x86_64.hpet", "Enable High Precision Event Timer clock", hpet_option);

struct hpet {
    struct clock clock;
    void* regs;
    uint32_t period;
};

constexpr uint32_t HPET_CAP = 0x0;
constexpr uint32_t HPET_CFG = 0x10;
constexpr uint32_t HPET_COUNTER = 0xF0;

static void hpet_reset(struct clock* c) {
    struct hpet* hpet = containerof(c, struct hpet, clock);
    mmio_write64(hpet->regs + HPET_COUNTER, 0);
}

static uint64_t hpet_get_elapsed_ns(struct clock* c) {
    struct hpet* hpet = containerof(c, struct hpet, clock);
    uint64_t counter = mmio_read64(hpet->regs + HPET_COUNTER);
    return (uint64_t)((__uint128_t)counter * (__uint128_t)hpet->period / 1'000'000);
}

static struct hpet hpet_clock = {
    .clock = {
        .name = "hpet",
        .priority = 128,
        .reset = hpet_reset,
        .get_elapsed_ns = hpet_get_elapsed_ns,
    },
};

static void hpet_init(struct boot_info*) {
    if (!use_hpet)
        return;

    uacpi_table table = {0};
    if (uacpi_table_find_by_signature("HPET", &table) != UACPI_STATUS_OK) {
        log("x86_64/hpet: No HPET table found\n");
        return;
    }

    struct acpi_hpet* hpet = table.ptr;

    uintptr_t vaddr = vm_alloc(0x1000);
    pmap_map(&kernel_vma.pmap, vaddr, hpet->address.address, PTE_READ | PTE_WRITE, CACHE_MMIO);
    hpet_clock.regs = (void*)vaddr;

    uacpi_table_unref(&table);
    log("x86_64/hpet: HPET registers mapped at %p\n", hpet_clock.regs);

    mmio_write64(hpet_clock.regs + HPET_CFG, mmio_read64(hpet_clock.regs + HPET_CFG) | 1);

    uint64_t caps = mmio_read64(hpet_clock.regs + HPET_CAP);
    hpet_clock.period = (caps >> 32) & 0xFFFF'FFFF;

    assert(clock_switch(&hpet_clock.clock), "Unable to switch to HPET!\n");
}
INIT_TASK(hpet_stage, "arch.x86_64.hpet", hpet_init);
INIT_DEPENDS(hpet_stage, acpi_stage);
INIT_ENTAILS(hpet_stage, percpu_stage);