// HAL APIC support. This file owns local APIC, I/O APIC, and interrupt-controller setup.

#include <kernel/hal/hal_apic.h>
#include <kernel/hal/hal_internal.h>
#include <kernel/hal/hal_api.h>
#include <kernel/hal/hal_timer.h>
#include <kernel/hal/hal_processor.h>
#include <kernel/ke_vectors.h>
#include <stdint.h>
#include <stdbool.h>
#include <intrin.h>


// ---- Local APIC ----.
static volatile uint32_t* gLapic = (volatile uint32_t*)0;
#define LAPIC_REG(off)   (*(volatile uint32_t*)((uint8_t*)gLapic + (off)))
#define LAPIC_ID         0x020
#define LAPIC_EOI        0x0B0
#define LAPIC_TPR        0x080
#define LAPIC_ISR_BASE   0x100
#define LAPIC_IRR_BASE   0x200
#define LAPIC_SPURIOUS   0x0F0
#define LAPIC_ICR_LO     0x300
#define LAPIC_ICR_HI     0x310
#define LAPIC_ENABLE     0x100   // bit in spurious vector register

#define IA32_APIC_BASE          0x0000001Bu
#define IA32_APIC_BASE_ENABLE   (1ULL << 11)
#define IA32_APIC_BASE_X2APIC   (1ULL << 10)
#define IA32_X2APIC_ID          0x00000802u
#define IA32_X2APIC_EOI         0x0000080Bu
#define IA32_X2APIC_TPR         0x00000808u
#define IA32_X2APIC_ISR_BASE    0x00000810u
#define IA32_X2APIC_IRR_BASE    0x00000820u
#define IA32_X2APIC_SVR         0x0000080Fu
#define IA32_X2APIC_ICR         0x00000830u

#define CPUID_FEATURE_ECX_X2APIC (1u << 21)


// ---- I/O APIC interrupt domains ----.
#define HAL_IOAPIC_MAX 8u
#define IOAPIC_IDX 0x00u
#define IOAPIC_DATA 0x10u

typedef struct HAL_IOAPIC {
    volatile uint32_t* Registers;
    uintptr_t Physical;
    uint32_t GsiBase;
    uint32_t Pins;
    volatile long RegisterLock;
} HAL_IOAPIC;

static HAL_IOAPIC gIoapics[HAL_IOAPIC_MAX];
static uint32_t gIoapicCount;
static uint32_t gBspLapicId;
static uint8_t gX2ApicSupported;
static uint8_t gX2ApicEnabled;
static uint32_t gIsaGsi[16];
static uint16_t gIsaFlags[16];
static bool gIsaOverridesInitialized;

// I/O APIC register selectors are shared by every CPU accessing that APIC.
static uint32_t IoapicRead(HAL_IOAPIC* apic, uint8_t reg) {
    uintptr_t interruptFlags = HalDisableInterrupts();
    while (_InterlockedCompareExchange(&apic->RegisterLock, 1, 0))
        _mm_pause();
    apic->Registers[IOAPIC_IDX / 4u] = reg;
    uint32_t value = apic->Registers[IOAPIC_DATA / 4u];
    _InterlockedExchange(&apic->RegisterLock, 0);
    HalRestoreInterrupts(interruptFlags);
    return value;
}

static void IoapicWrite(HAL_IOAPIC* apic, uint8_t reg, uint32_t value) {
    uintptr_t interruptFlags = HalDisableInterrupts();
    while (_InterlockedCompareExchange(&apic->RegisterLock, 1, 0))
        _mm_pause();
    apic->Registers[IOAPIC_IDX / 4u] = reg;
    apic->Registers[IOAPIC_DATA / 4u] = value;
    _InterlockedExchange(&apic->RegisterLock, 0);
    HalRestoreInterrupts(interruptFlags);
}

static HAL_IOAPIC* HalApicFindGsi(uint32_t gsi, uint32_t* pinOut) {
    for (uint32_t i = 0; i < gIoapicCount; i++) {
        HAL_IOAPIC* apic = &gIoapics[i];
        if (gsi >= apic->GsiBase && gsi - apic->GsiBase < apic->Pins) {
            *pinOut = gsi - apic->GsiBase;
            return apic;
        }
    }
    return (HAL_IOAPIC*)0;
}

static void HalApicDetectMode(void) {
    int regs[4] = {0, 0, 0, 0};
    __cpuid(regs, 1);
    gX2ApicSupported = (regs[2] & CPUID_FEATURE_ECX_X2APIC) ? 1u : 0u;

    uint64_t apicBase = __readmsr(IA32_APIC_BASE);
    gX2ApicEnabled = (gX2ApicSupported && (apicBase & IA32_APIC_BASE_X2APIC)) ? 1u : 0u;
}

// Stalls for the requested bootstrap interval using the HAL timer fallback chain.
static void HalApicStallUs(uint32_t us) {
    HalTimerStall(us);
}

static void HalApicWaitIcrIdle(void) {
    if (gX2ApicEnabled)
        return;

    for (uint32_t i = 0; i < 1000000u; i++) {
        if (!(LAPIC_REG(LAPIC_ICR_LO) & (1u << 12)))
            return;
        _mm_pause();
    }
}

static void HalApicSendIpi(uint32_t lapicId, uint32_t low) {
    if (gX2ApicEnabled) {
        __writemsr(IA32_X2APIC_ICR, ((uint64_t)lapicId << 32) | low);
        return;
    }

    HalApicWaitIcrIdle();
    LAPIC_REG(LAPIC_ICR_HI) = lapicId << 24;
    LAPIC_REG(LAPIC_ICR_LO) = low;
    HalApicWaitIcrIdle();
}

// IRQ n -> redirection table at register 0x10 + 2*n (low), 0x11 + 2*n (high).
#define IOREDTBL_LO(n)  (0x10 + 2*(n))
#define IOREDTBL_HI(n)  (0x11 + 2*(n))

// Returns the local interrupt-controller hardware ID of the boot processor.
uint32_t HalApicGetBootProcessorId(void) {
    return gBspLapicId;
}

// Returns whether the local interrupt controller is using x2APIC mode.
uint8_t HalApicUsingX2Apic(void) {
    return gX2ApicEnabled;
}

// Returns whether the active local interrupt-controller mode can address one hardware ID.
uint8_t HalApicCanAddressId(uint32_t hardwareId) {
    if (gX2ApicEnabled)
        return 1u;
    return (hardwareId < 256u) ? 1u : 0u;
}

// Initializes the local interrupt controller on the current processor.
void HalApicInitializeCurrentProcessor(void) {
    uint64_t apicBase = __readmsr(IA32_APIC_BASE) | IA32_APIC_BASE_ENABLE;

    if (gX2ApicEnabled) {
        apicBase |= IA32_APIC_BASE_X2APIC;
        __writemsr(IA32_APIC_BASE, apicBase);
        __writemsr(IA32_X2APIC_SVR, LAPIC_ENABLE | 0xFFu);
        HalApicSetTaskPriority(0);
        return;
    }

    __writemsr(IA32_APIC_BASE, apicBase);
    if (gLapic)
        LAPIC_REG(LAPIC_SPURIOUS) = LAPIC_ENABLE | 0xFF;
    HalApicSetTaskPriority(0);
}

// Controls the current CPU's interrupt priority without changing its interrupt-enable flag.
void HalApicSetTaskPriority(uint8_t priority) {
    if (gX2ApicEnabled) {
        __writemsr(IA32_X2APIC_TPR, priority);
    } else if (gLapic) {
        LAPIC_REG(LAPIC_TPR) = priority;
    }
}

uint8_t HalApicQueryTaskPriority(void) {
    if (gX2ApicEnabled)
        return (uint8_t)__readmsr(IA32_X2APIC_TPR);
    return gLapic ? (uint8_t)LAPIC_REG(LAPIC_TPR) : 0u;
}

// Reads the current processor's local delivery state after a source is masked.
bool HalApicVectorPending(uint8_t vector) {
    const uint32_t index = (uint32_t)vector >> 5u;
    const uint32_t mask = 1u << (vector & 31u);
    uint32_t irr, isr;
    if (gX2ApicEnabled) {
        irr = (uint32_t)__readmsr(IA32_X2APIC_IRR_BASE + index);
        isr = (uint32_t)__readmsr(IA32_X2APIC_ISR_BASE + index);
    } else if (gLapic) {
        irr = LAPIC_REG(LAPIC_IRR_BASE + 0x10u * index);
        isr = LAPIC_REG(LAPIC_ISR_BASE + 0x10u * index);
    } else {
        return true;
    }
    return ((irr | isr) & mask) != 0u;
}

// Broadcasts an NMI to every processor except the caller.
void HalApicBroadcastNmi(void) {
    if (!gLapic && !gX2ApicEnabled) return;

    // Stop-the-world panic IPI. Use NMI delivery instead of a normal fixed
    // vector so CPUs running with IF=0 inside spinlock/high-IRQL regions still
    // enter the panic path. The NMI vector is already installed in the IDT on
    // every CPU that has reached kernel-owned state.
    HalApicSendIpi(0, (3u << 18) | (4u << 8)); // all-excluding-self, NMI
}

// Sends one fixed-vector processor interrupt without exposing APIC mechanics.
void HalApicSendProcessorInterrupt(uint32_t hardwareId, uint8_t vector) {
    if ((!gLapic && !gX2ApicEnabled) || vector < 0x20u || !HalApicCanAddressId(hardwareId))
        return;
    HalApicSendIpi(hardwareId, vector);
}

// Sends one non-maskable processor interrupt.
void HalApicSendProcessorNmi(uint32_t hardwareId) {
    if ((!gLapic && !gX2ApicEnabled) || !HalApicCanAddressId(hardwareId))
        return;
    HalApicSendIpi(hardwareId, (4u << 8));
}

// Starts one secondary processor through the local interrupt-controller startup protocol.
KSTATUS HalApicStartProcessor(uint32_t hardwareId, uintptr_t startPhysical) {
    if ((!gLapic && !gX2ApicEnabled) || !startPhysical ||
        (startPhysical & 0xFFFu) != 0 || startPhysical >= 0x100000u ||
        hardwareId == gBspLapicId) {
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (!HalApicCanAddressId(hardwareId))
        return NTLSTATUS_NOT_SUPPORTED;

    uint8_t startupVector = (uint8_t)(startPhysical >> 12);
    // INIT assert/deassert, then two SIPIs. The target begins at startPhysical.
    HalApicSendIpi(hardwareId, 0x0000C500u); // INIT, level assert, level-triggered
    HalApicStallUs(10000);
    HalApicSendIpi(hardwareId, 0x00008500u); // INIT, level deassert
    HalApicStallUs(200);
    HalApicSendIpi(hardwareId, 0x00000600u | startupVector); // SIPI
    HalApicStallUs(200);
    HalApicSendIpi(hardwareId, 0x00000600u | startupVector); // SIPI retry
    HalApicStallUs(200);
    return NTLSTATUS_SUCCESS;
}

// MSI destination encoding without interrupt remapping supports an 8-bit APIC ID.
KSTATUS HalApicComposeMsi(uint32_t hardwareId, uint8_t vector, HAL_MSI_MESSAGE* out) {
    if (!out || vector < KE_VECTOR_DEVICE_FIRST || vector > KE_VECTOR_DEVICE_LAST)
        return NTLSTATUS_INVALID_PARAMETER;
    if (hardwareId > 255u)
        return NTLSTATUS_NOT_SUPPORTED;
    out->Address = 0xFEE00000ull | ((uint64_t)hardwareId << 12);
    out->Data = vector;
    out->Reserved = 0;
    return NTLSTATUS_SUCCESS;
}

static void HalApicInitializeIsaMap(void) {
    if (gIsaOverridesInitialized) return;
    for (uint32_t irq = 0; irq < 16; irq++) {
        gIsaGsi[irq] = irq;
        gIsaFlags[irq] = 0;
    }
    gIsaOverridesInitialized = true;
}

void HalApicAddIoapic(uintptr_t physical, uint32_t gsiBase) {
    if (!physical || gIoapicCount >= HAL_IOAPIC_MAX)
        return;
    for (uint32_t i = 0; i < gIoapicCount; i++)
        if (gIoapics[i].Physical == physical)
            return;
    HAL_IOAPIC* apic = &gIoapics[gIoapicCount++];
    apic->Physical = physical;
    apic->GsiBase = gsiBase;
}

void HalApicSetIsaOverride(uint8_t irq, uint32_t gsi, uint16_t flags) {
    if (irq >= 16u) return;
    HalApicInitializeIsaMap();
    gIsaGsi[irq] = gsi;
    gIsaFlags[irq] = flags;
}

KSTATUS HalApicResolveIsa(uint8_t irq, uint32_t* gsiOut, uint16_t* flagsOut) {
    if (irq >= 16u || !gsiOut || !flagsOut)
        return NTLSTATUS_INVALID_PARAMETER;
    HalApicInitializeIsaMap();
    uint32_t pin;
    if (!HalApicFindGsi(gIsaGsi[irq], &pin))
        return NTLSTATUS_NOT_SUPPORTED;
    *gsiOut = gIsaGsi[irq];
    *flagsOut = gIsaFlags[irq];
    return NTLSTATUS_SUCCESS;
}

// ACPI polarity bits 1:0 and trigger bits 3:2 encode bus-defined defaults.
KSTATUS HalApicRouteGsi(uint32_t gsi, uint8_t vector, uint32_t hardwareId,
                       uint16_t flags, bool masked) {
    if (vector < KE_VECTOR_DEVICE_FIRST || vector > KE_VECTOR_DEVICE_LAST || hardwareId > 255u)
        return NTLSTATUS_INVALID_PARAMETER;
    uint32_t pin;
    HAL_IOAPIC* apic = HalApicFindGsi(gsi, &pin);
    if (!apic) return NTLSTATUS_NOT_SUPPORTED;
    uint32_t mode = 0;
    if ((flags & 3u) == 3u) mode |= (1u << 13);
    if (((flags >> 2) & 3u) == 3u) mode |= (1u << 15);
    IoapicWrite(apic, IOREDTBL_LO(pin), (1u << 16) | vector | mode);
    IoapicWrite(apic, IOREDTBL_HI(pin), (hardwareId & 255u) << 24);
    if (!masked) IoapicWrite(apic, IOREDTBL_LO(pin), vector | mode);
    return NTLSTATUS_SUCCESS;
}

KSTATUS HalApicMaskGsi(uint32_t gsi, bool masked) {
    uint32_t pin;
    HAL_IOAPIC* apic = HalApicFindGsi(gsi, &pin);
    if (!apic) return NTLSTATUS_NOT_SUPPORTED;
    uint32_t lo = IoapicRead(apic, IOREDTBL_LO(pin));
    if (masked) lo |= 1u << 16;
    else lo &= ~(1u << 16);
    IoapicWrite(apic, IOREDTBL_LO(pin), lo);
    return NTLSTATUS_SUCCESS;
}

void HalApicSendEoi(void) {
    if (gX2ApicEnabled) {
        __writemsr(IA32_X2APIC_EOI, 0);
        return;
    }
    LAPIC_REG(LAPIC_EOI) = 0;
}

void HalApicInitialize(uintptr_t lapicPhys) {
    // The PIC remains masked: legacy input arrives through programmed I/O APICs.
    __outbyte(0x21, 0xFF);
    __outbyte(0xA1, 0xFF);
    HalApicDetectMode();

    gLapic = (volatile uint32_t*)HalPhysicalToVirtual(lapicPhys);
    if (gX2ApicEnabled)
        gBspLapicId = (uint32_t)__readmsr(IA32_X2APIC_ID);
    else
        gBspLapicId = (LAPIC_REG(LAPIC_ID) >> 24) & 0xFFu;
    HalApicInitializeCurrentProcessor();
    if (gIoapicCount == 0)
        HalApicAddIoapic(0xFEC00000u, 0);

    for (uint32_t index = 0; index < gIoapicCount; index++) {
        HAL_IOAPIC* apic = &gIoapics[index];
        apic->Registers = (volatile uint32_t*)HalPhysicalToVirtual(apic->Physical);
        uint32_t version = IoapicRead(apic, 1u);
        apic->Pins = ((version >> 16) & 0xFFu) + 1u;
        for (uint32_t pin = 0; pin < apic->Pins; pin++) {
            IoapicWrite(apic, IOREDTBL_LO(pin), (1u << 16) | KE_VECTOR_DEVICE_FIRST);
            IoapicWrite(apic, IOREDTBL_HI(pin), (gBspLapicId & 255u) << 24);
        }
    }
}