// hal_timer.c - Local APIC timer, calibrated against PIT channel 2
#include "hal_api.h"
#include "hal_apic.h"
#include "hal_debug.h"
#include <stdint.h>
#include <intrin.h>

#define LAPIC_TIMER_LVT   0x320
#define LAPIC_TIMER_INIT  0x380
#define LAPIC_TIMER_CUR   0x390
#define LAPIC_TIMER_DIV   0x3E0

#define TIMER_VECTOR      0x20      // vector 0x20 = IRQ 0 slot in IDT
#define TIMER_PERIODIC    (1 << 17)
#define TIMER_MASKED      (1 << 16)
#define TIMER_DIV16       0x3

#define QUANTUM_MS        10
#define PIT_HZ            1193182UL

extern HAL_API gApi;
static volatile uint32_t* gLapic = NULL;
#define LAPIC_REG(off) (*(volatile uint32_t *)((uint8_t *)gLapic + (off)))

static volatile uint64_t gTicks = 0;
void (*gTimerCallback)(void) = NULL;

static inline void outb(uint16_t port, uint8_t val) { __outbyte(port, val); }
static inline uint8_t inb(uint16_t port)            { return __inbyte(port); }

static void TimerIsr(void) {
    gTicks++;
    if (gTimerCallback) gTimerCallback();
    // EOI is sent by KeIrqDispatch after this returns
}

static uint64_t QueryTicks(void) { return gTicks; }

static void Stall(uint32_t us) {
    uint64_t start = gTicks;
    uint64_t ms    = (us + 999) / 1000;
    while ((gTicks - start) < ms) _mm_pause();
}

// Returns HPET ticks/second, or 0 if HPET is unavailable.
static uint64_t HpetInit(uintptr_t hpetPhys) {
    // HPET register block, accessed through physical window
    volatile uint64_t *hpet = (volatile uint64_t *)gApi.PhysToVirt(hpetPhys);

    // General Capabilities register (offset 0):
    // Bits 63:32 = counter clock period in femtoseconds
    uint64_t caps = hpet[0];
    uint32_t periodFs = (uint32_t)(caps >> 32);
    if (!periodFs) return 0;

    // Enable the HPET main counter (General Configuration, offset 0x10, bit 0)
    hpet[2] |= 1;   // offset 0x10 / 8 = index 2

    return (uint64_t)1e15 / periodFs;   // ticks per second
}

// Use PIT channel 2 one-shot to measure how many APIC ticks elapse in 10 ms, then return ticks/ms.
static uint32_t CalibrateApicTimerViaPit(void) {
    // Divide by 16
    LAPIC_REG(LAPIC_TIMER_DIV)  = TIMER_DIV16;
    // Mask, set vector - leave masked during calibration
    LAPIC_REG(LAPIC_TIMER_LVT)  = TIMER_MASKED | TIMER_VECTOR;
    // Load max count
    LAPIC_REG(LAPIC_TIMER_INIT) = 0xFFFFFFFF;

    // PIT channel 2: one-shot, ~10 ms
    uint16_t pit_count = (uint16_t)(PIT_HZ / 100); // 11931 ≈ 10 ms
    outb(0x61, (inb(0x61) & ~0x02) | 0x01);        // gate on, speaker off
    outb(0x43, 0xB0);                              // ch2, lo/hi, mode 0, binary
    outb(0x42, (uint8_t)(pit_count & 0xFF));
    outb(0x42, (uint8_t)(pit_count >> 8));

    // Wait for OUT2 (bit 5 of port 0x61) to go high
    while (!(inb(0x61) & 0x20)) _mm_pause();

    uint32_t elapsed = 0xFFFFFFFF - LAPIC_REG(LAPIC_TIMER_CUR);
    return elapsed / 10;   // ticks per ms
}

// Use HPET main counter to measure how many APIC ticks elapse in 10 ms, then return ticks/ms.
static uint32_t CalibrateApicTimerViaHpet(uintptr_t hpetPhys) {
    volatile uint64_t *hpet = (volatile uint64_t *)gApi.PhysToVirt(hpetPhys);
    uint64_t hpetHz = HpetInit(hpetPhys);
    if (!hpetHz) return 0;

    // Ticks in 10 ms = hpetHz / 100
    uint64_t ticksFor10ms = hpetHz / 100;

    LAPIC_REG(LAPIC_TIMER_DIV)  = TIMER_DIV16;
    LAPIC_REG(LAPIC_TIMER_LVT)  = TIMER_MASKED | TIMER_VECTOR;
    LAPIC_REG(LAPIC_TIMER_INIT) = 0xFFFFFFFF;

    // HPET main counter is at offset 0xF0, index 0xF0/8 = 30
    uint64_t hpetStart = hpet[30];
    while ((hpet[30] - hpetStart) < ticksFor10ms) _mm_pause();

    uint32_t elapsed = 0xFFFFFFFF - LAPIC_REG(LAPIC_TIMER_CUR);
    return elapsed / 10;   // APIC ticks per ms
}

KSTATUS HalTimerInitialize(uintptr_t lapicPhys, uintptr_t hpetPhys, HAL_TIMER *out) {
    gLapic = (volatile uint32_t *)gApi.PhysToVirt(lapicPhys);
    HalDbgPrint("HAL: lapicPhys=0x%llX, lapicVirt=0x%llX, hpetPhys=0x%llX\n", lapicPhys, (uintptr_t)gLapic, hpetPhys);

    uint32_t ticksPerMs = 0;
    if (hpetPhys) {
        HalDbgPrint("HAL: Attempting to calibrate APIC timer via HPET...\n");
        ticksPerMs = CalibrateApicTimerViaHpet(hpetPhys);
    }
    if (!ticksPerMs) {
        HalDbgPrint("HAL: HPET unavailable or calibration failed. Falling back to PIT...\n");
        ticksPerMs = CalibrateApicTimerViaPit();
    }
    if (!ticksPerMs) {
        HalDbgPrint("HAL: APIC calibration failed. How?\n");
        return STATUS_DEVICE_ERROR;
    }

    HalDbgPrint("HAL: APIC timer: %u ticks/ms\n", ticksPerMs);

    // Configure LAPIC timer in periodic mode
    LAPIC_REG(LAPIC_TIMER_DIV)  = TIMER_DIV16;
    LAPIC_REG(LAPIC_TIMER_LVT)  = TIMER_PERIODIC | TIMER_VECTOR;
    LAPIC_REG(LAPIC_TIMER_INIT) = ticksPerMs * QUANTUM_MS;

    out->QueryTicks = QueryTicks;
    out->Stall      = Stall;
    out->Isr        = TimerIsr;
    out->IrqNumber  = 0;
    return STATUS_SUCCESS;
}