/*
 * 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 <machine/idt.h>
#include <hal/mmio.h>
#include <hal/prim.h>
#include <drivers/acpi/tables.h>
#include <drivers/acpi/acpi.h>
#include <mm/vm.h>
#include <ex/trace.h>
#include <ke/bugcheck.h>

/* Local APIC timer vector */
#define LAPIC_TMR_VECTOR 0x81

/*
 * 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);

/* Globals */
extern VOID LapicTmrIsr(VOID);

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

/*
 * Start the Local APIC timer
 *
 * @Mcb:    Machine core block
 * @Mask:   True to mask the timer
 * @Mode:   Timer mode
 * @Count:  Timer count
 */
static VOID
LapicTimerStart(MCB *Mcb, BOOLEAN Mask, UCHAR Mode, ULONG Count)
{
    ULONG Tmp;

    Tmp = (Mode << 17) | (Mask << 16) | LAPIC_TMR_VECTOR;
    LapicRegWrite(Mcb, LAPIC_LVT_TMR, Tmp);
    LapicRegWrite(Mcb, LAPIC_DCR, 0);
    LapicRegWrite(Mcb, LAPIC_INIT_CNT, Count);
}

/*
 * Start the Local APIC timer in oneshot mode
 *
 * @Mcb:    Machine core block
 * @Mask:   True to mask the timer
 * @Count:  Initial count
 */
static void
LapicTimerOneShot(MCB *Mcb, BOOLEAN Mask, ULONG Count)
{
    LapicTimerStart(Mcb, Mask, LVT_TMR_ONESHOT, Count);
}

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

ULONG
MdLapicId(VOID)
{
    KPCR *ThisCore;
    MCB *Mcb;

    ThisCore = HalKpcrCurrent();
    Mcb = &ThisCore->Mcb;

    /* 32-bit when in x2APIC mode */
    if (Mcb->HasX2Apic) {
        return LapicRegRead(Mcb, LAPIC_ID) & 0xFFFFFFFF;
    }

    return (LapicRegRead(Mcb, LAPIC_ID) >> 24) & 0xF;
}

VOID
MdLapicTimerOneshotUs(USIZE Usec)
{
    UQUAD Ticks;
    KPCR *Kpcr;
    MCB *Mcb;

    Kpcr = HalKpcrCurrent();
    Mcb = &Kpcr->Mcb;

    Ticks = Usec * (Mcb->LapicTmrFreq / 1000000);
    LapicTimerOneShot(Mcb, false, Ticks);
}

VOID
MdLapicSendEoi(VOID)
{
    KPCR *Kpcr;

    Kpcr = HalKpcrCurrent();
    LapicRegWrite(&Kpcr->Mcb, LAPIC_EOI, 0);
}

VOID
MdLapicInit(KPCR *Kpcr)
{
    ACPI_MADT *Madt;
    MCB *Mcb;

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

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

    /* Obtain the ACPI MADT table */
    Madt = AcpiQuery("APIC");
    if (Madt == NULL) {
        KeBugCheck(BUGCHECK_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);

    /* Set the LAPIC timer ISR gate */
    MdIdtSetGate(
        0x81,
        (UPTR)LapicTmrIsr,
        IDT_INT_GATE,
        0
    );
}
