// hal_timer.c - Local APIC timer
#include <stdint.h>
#include <intrin.h>
#include <hal_api.h>
#include <hal_apic.h>
#include <hal_debug.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;
static uint32_t gApicTicksPerUs = 0;
static uint32_t gApicTicksPerQuantum = 0;
static uint64_t gHpetTicksPerUs = 0;
static uint64_t gHpetHz = 0;
static volatile uint64_t* gHpet = NULL;
#define LAPIC_REG(off) (*(volatile uint32_t *)((uint8_t *)gLapic + (off)))

static volatile uint64_t gTicks = 0;
void (* volatile 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) {
    // Prefer HPET when present: it is already a monotonically increasing main
    // counter and avoids races between the LAPIC periodic tick and the current
    // countdown register.
    if (gHpet && gHpetHz) {
        uint64_t counter = gHpet[30]; // offset 0xF0: main counter
        uint64_t seconds = counter / gHpetHz;
        uint64_t rem = counter % gHpetHz;
        return seconds * 1000000ull + (rem * 1000000ull) / gHpetHz;
    }

    // Fallback: combine the coarse periodic tick with the current LAPIC
    // countdown.  Read the coarse tick twice so an interrupt between gTicks and
    // LAPIC_TIMER_CUR cannot make time jump backwards by almost one quantum.
    uint64_t baseTicks1;
    uint64_t baseTicks2;
    uint32_t current;
    do {
        baseTicks1 = gTicks;
        current = LAPIC_REG(LAPIC_TIMER_CUR);
        baseTicks2 = gTicks;
    } while (baseTicks1 != baseTicks2);

    uint32_t initial = gApicTicksPerQuantum;
    
    uint32_t elapsedApicTicks = 0;
    if (current < initial) {
        elapsedApicTicks = initial - current;
    }

    uint64_t subQuantumUs = 0;
    if (gApicTicksPerUs)
        subQuantumUs = elapsedApicTicks / gApicTicksPerUs;

    return (baseTicks1 * QUANTUM_MS * 1000ull) + subQuantumUs;
}

static void Stall(uint32_t us) {
    // Priority 1: Use HPET (Highest Precision)
    if (gHpet && gHpetTicksPerUs > 0) {
        uint64_t start = gHpet[30]; // Offset 0xF0 is the Main Counter
        uint64_t delta = (uint64_t)us * gHpetTicksPerUs;
        while ((gHpet[30] - start) < delta) {
            _mm_pause(); // Hint to CPU we are in a spin-loop
        }
        return;
    }

    // Priority 2: Use APIC Timer (High Precision, but handles wrap-around)
    if (gApicTicksPerUs > 0) {
        uint32_t total_ticks = us * gApicTicksPerUs;
        uint32_t last_val = LAPIC_REG(LAPIC_TIMER_CUR);
        uint32_t elapsed = 0;

        while (elapsed < total_ticks) {
            uint32_t current = LAPIC_REG(LAPIC_TIMER_CUR);
            if (current < last_val) {
                elapsed += (last_val - current);
            } else if (current > last_val) {
                // Handled the reset/wrap at the end of QUANTUM_MS
                uint32_t init_val = LAPIC_REG(LAPIC_TIMER_INIT);
                elapsed += (last_val + (init_val - current));
            }
            last_val = current;
            _mm_pause();
        }
        return;
    }

    // Priority 3: Fallback to Port 0x80
    for (uint32_t i = 0; i < us; i++) {
        __outbyte(0x80, 0);
    }
}

// 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;
    }

    gApicTicksPerUs = ticksPerMs / 1000;
    gApicTicksPerQuantum = ticksPerMs * QUANTUM_MS;

    if (hpetPhys) {
        gHpet = (volatile uint64_t*)gApi.PhysToVirt(hpetPhys);
        uint64_t hpetHz = HpetInit(hpetPhys);
        gHpetHz = hpetHz;
        gHpetTicksPerUs = hpetHz / 1000000;
    }

    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) = gApicTicksPerQuantum;

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