#include <zinnia/clock.h>
#include <zinnia/cmdline.h>
#include <zinnia/mmio.h>
#include <zinnia/print.h>
#include <zinnia/utils.h>
#include <assert.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("hpet", hpet_option);

struct hpet {
    struct clock clock;
    struct mmio_view regs;
    uint32_t period;
};

constexpr mmio_reg32 HPET_CAP = {.offset = 0};
constexpr mmio_reg32 HPET_CFG = {.offset = 0x10};
constexpr mmio_reg32 HPET_COUNTER = {.offset = 0xF0};

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

static uint64_t hpet_get_elapsed_ns(struct clock* c) {
    struct hpet* hpet = containerof(c, struct hpet, clock);
    uint64_t counter = mmio_load(&hpet->regs, HPET_COUNTER);
    return (counter * (uint64_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 char table_buf[0x1000];

void hpet_init() {
    if (!use_hpet)
        return;

    uacpi_setup_early_table_access(table_buf, sizeof(table_buf));
    uacpi_table table = {0};
    assert(!uacpi_table_find_by_signature("HPET", &table), "HPET table not found!\n");

    struct acpi_hpet* hpet = table.ptr;
    hpet_clock.regs = mmio_new(hpet->address.address, 0x1000);
    uacpi_table_unref(&table);
    printf("HPET regs at %p\n", hpet_clock.regs.base);

    mmio_store(&hpet_clock.regs, HPET_CFG, mmio_load(&hpet_clock.regs, HPET_CFG) | 1);

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

    assert(clock_switch(&hpet_clock.clock), "Unable to switch to HPET!\n");
}