/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Local APIC driver
 * Author:      Chloe M.
 */

#include <machine/lapic.h>
#include <machine/lapicreg.h>
#include <machine/msr.h>
#include <machine/cpuid.h>
#include <machine/i8254.h>
#include <hal/mmio.h>
#include <hal/prim.h>
#include <drivers/acpi/tables.h>
#include <drivers/acpi/acpi.h>
#include <mm/vmm.h>
#include <ex/trace.h>
#include <ke/knot.h>
#include <stdef.h>

/*
 * Trace only on the bootstrap processor but not on any others
 * to avoid log spam...
 */
#define DTRACE_BSP(Fmt, ...) do {                       \
        ULONG ApicBase;                                 \
                                                        \
        ApicBase = MdRdmsr(IA32_APIC_BASE_MSR);         \
        if (ISSET(ApicBase, BIT(8))) {                  \
            TRACE("[ LAPIC ]: " Fmt, ##__VA_ARGS__);    \
        }                                               \
    } while (0);

/*
 * Returns true if the Local APIC unit is supported on the
 * current processor.
 */
static BOOLEAN
LapicIsSupported(VOID)
{
    ULONG Edx, Unused;

    CPUID(1, Unused, Unused, Unused, Edx);
    return ISSET(Edx, BIT(9)) != 0;
}

/*
 * Returns true if the CPU can operate its Local APIC unit
 * in the newer x2APIC mode.
 */
static BOOLEAN
LapicHasX2Apic(VOID)
{
    ULONG Ecx, Unused;

    CPUID(1, Unused, Unused, Ecx, Unused);
    return ISSET(Ecx, BIT(21)) != 0;
}

/*
 * Read a Local APIC register
 *
 * @Mcb:        Machine core block
 * @Register:   Register to read
 */
static UQUAD
LapicRegRead(MCB *Mcb, ULONG Register)
{
    ULONG *Base;
    UQUAD Value;

    if (!Mcb->HasX2Apic) {
        Base = PTR_OFFSET(Mcb->LapicBase, Register);
        Value = MMIORead32(Base);
    } else {
        Register >>= 4;
        Value = MdRdmsr(x2APIC_MSR_BASE + Register);
    }

    return Value;
}

/*
 * Write a value to a Local APIC register space
 *
 * @Mcb:        Machine core block
 * @Register:   Register to read
 * @Value:      Value to write
 */
static VOID
LapicRegWrite(MCB *Mcb, ULONG Register, UQUAD Value)
{
    ULONG *Base;

    if (!Mcb->HasX2Apic) {
        Base = PTR_OFFSET(Mcb->LapicBase, Register);
        MMIOWrite32(Base, (ULONG)Value);
    } else {
        Register >>= 4;
        MdWrmsr(x2APIC_MSR_BASE + Register, Value);
    }
}

/*
 * Enable the Local APIC unit
 *
 * @Mcb:   Machine core block
 */
static VOID
LapicEnable(MCB *Mcb)
{
    UQUAD ApicBase;
    ULONG Svr, Version;
    const CHAR *mode;
    const CHAR *type = "discrete 82489DX";

    /* Hardware enable the Local APIC unit */
    ApicBase = MdRdmsr(IA32_APIC_BASE_MSR);
    ApicBase |= LAPIC_HW_ENABLE;
    ApicBase |= Mcb->HasX2Apic << x2APIC_ENABLE_SHIFT;
    MdWrmsr(IA32_APIC_BASE_MSR, ApicBase);

    /* Software enable the Local APIC unit */
    Svr = LapicRegRead(Mcb, LAPIC_SVR);
    Svr |= LAPIC_SW_ENABLE;
    LapicRegWrite(Mcb, LAPIC_SVR, Svr);

    /* Check the type */
    Version = LapicRegRead(Mcb, LAPIC_VERSION);
    if ((Version & 0xFF) > 0) {
        type = "integrated";
    }

    mode = Mcb->HasX2Apic ? "x2apic" : "xapic";
    DTRACE_BSP("enabled %s %s\n", type, mode);
}

/*
 * Stop the Local APIC timer
 *
 * @Mcb: Machine core block
 */
static VOID
LapicTimerStop(MCB *Mcb)
{
    if (Mcb == NULL) {
        return;
    }

    LapicRegWrite(Mcb, LAPIC_LVT_TMR, LAPIC_LVT_MASK);
    LapicRegWrite(Mcb, LAPIC_INIT_CNT, 0);
}

/*
 * Calibrate the Local APIC timer
 *
 * @Mcb: Machine core block
 */
static VOID
LapicTimerCalibrate(MCB *Mcb)
{
    const USHORT MAX_SAMPLES = 0xFFFF;
    USHORT TicksStart, TicksEnd;
    USIZE Freq, TicksTotal;

    if (Mcb == NULL) {
        return;
    }

    LapicTimerStop(Mcb);
    MdPitSetCount(MAX_SAMPLES);
    TicksStart = MdPitGetCount();

    LapicRegWrite(Mcb, LAPIC_INIT_CNT, MAX_SAMPLES);
    while (LapicRegRead(Mcb, LAPIC_CUR_CNT) != 0);

    TicksEnd = MdPitGetCount();
    TicksTotal = TicksStart - TicksEnd;

    Freq = (MAX_SAMPLES / TicksTotal) * I8254_DIVIDEND;
    LapicTimerStop(Mcb);
    Mcb->LapicTmrFreq = Freq;
}

VOID
MdLapicSendIpi(UCHAR Vector, UCHAR DestId, BOOLEAN LogicalDest,
    IPI_SHORTHAND Xnd, IPI_DELMOD DelMod)
{
    KPCR *ThisCore;
    MCB *Mcb;
    ULONG IcrLow, IcrHigh;

    /*
     * If the shorthand is refering to ourselves and we are x2APIC we
     * can use a better path via its specialized SELF IPI MSR.
     */
    ThisCore = HalKpcrCurrent();
    if (Xnd == IPI_XND_SELF && Mcb->HasX2Apic) {
        MdWrmsr(LAPIC_SELF_IPI, Vector);
        return;
    }

    Mcb = &ThisCore->Mcb;
    IcrLow = 0;
    IcrHigh = 0;

    /* Clamp these fields */
    DelMod &= 7;
    Xnd &= 3;

    /* Encode the ICR */
    IcrLow |= Vector & 0xFF;
    IcrLow |= Xnd << 18;
    IcrLow |= DelMod << 8;
    IcrLow |= LogicalDest << 11;

    /* For xAPIC only */
    if (!Mcb->HasX2Apic) {
        IcrHigh |= (UQUAD)DestId << 24;
        LapicRegWrite(Mcb, LAPIC_ICRHI, IcrHigh);
        LapicRegWrite(Mcb, LAPIC_ICRLO, IcrLow);

        /* Only on legacy xAPICs do we poll */
        for (;;) {
            IcrLow = LapicRegRead(Mcb, LAPIC_ICRLO);
            if (ISSET(IcrLow, IPI_DELSTAT_PENDING))
                HalCpuSpinWait();
        }

        return;
    }

    /* On x2APICs only as it queues */
    LapicRegWrite(Mcb, LAPIC_ICRLO, ((UQUAD)DestId << 32) | IcrLow);
}

VOID
MdLapicInit(KPCR *Kpcr)
{
    INTR_HANDLER Handler;
    INTR_DATA *IntrData;
    ACPI_MADT *Madt;
    MCB *Mcb;

    if (Kpcr == NULL) {
        KeKnot(
            KNOT_MISC,
            "lapic: failed to initialize lapic driver\n"
        );
    }

    if (!LapicIsSupported()) {
        KeKnot(
            KNOT_MISSING_HARDWARE,
            "lapic: local apic not supported\n"
        );
    }

    /* Obtain the ACPI MADT table */
    Madt = AcpiQuery("APIC");
    if (Madt == NULL) {
        KeKnot(KNOT_MISC, "lapic: could not obtain MADT\n");
    }

    Mcb = &Kpcr->Mcb;
    Mcb->LapicBase = PMA_TO_VMA((UPTR)Madt->LapicAddr);
    Mcb->HasX2Apic = LapicHasX2Apic();

    /* Enable the Local APIC */
    LapicEnable(Mcb);

    /* Calibrate the timer */
    LapicTimerCalibrate(Mcb);
}