// CPU-local kernel state. This file owns bootstrap and per-CPU bookkeeping shared by scheduler and SMP code.

#include <ntlite/atomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <intrin.h>
#include <kernel/ke_hal.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_debug.h>
#include <kernel/hal/hal_api.h>

#define BOOTSTRAP_CPU_STACK_SIZE 0x4000u
#define IA32_GS_BASE             0xC0000101u
#define IA32_KERNEL_GS_BASE      0xC0000102u
#define IA32_MTRRCAP             0x000000FEu
#define IA32_MTRR_PHYSBASE0      0x00000200u
#define IA32_MTRR_PHYSMASK0      0x00000201u
#define IA32_PAT                 0x00000277u
#define IA32_MTRR_DEF_TYPE       0x000002FFu
#define CPUID_FEATURE_MTRR       (1u << 12)
#define CPUID_FEATURE_PAT        (1u << 16)
#define MTRR_CAP_VARIABLE_COUNT  0x000000FFull
#define MTRR_MASK_VALID          (1ull << 11)
#define MTRR_DEF_TYPE_MASK       0x000000FFull
#define MTRR_DEF_FIXED_ENABLE    (1ull << 10)
#define MTRR_DEF_ENABLE          (1ull << 11)
#define MEMORY_TYPE_UC           0x00ull
#define MEMORY_TYPE_WC           0x01ull
#define MEMORY_TYPE_WB           0x06ull
#define CR0_CACHE_DISABLE        (1ull << 30)
#define CR0_NOT_WRITE_THROUGH    (1ull << 29)
#define CR4_PAGE_GLOBAL_ENABLE   (1ull << 7)
#define CR4_PCID_ENABLE          (1ull << 17)
#define PAT_ENTRY(type, index)   ((uint64_t)(type) << ((index) * 8u))
#define CANONICAL_PAT            (PAT_ENTRY(MEMORY_TYPE_WB, 0) | \
                                  PAT_ENTRY(MEMORY_TYPE_WC, 1) | \
                                  PAT_ENTRY(MEMORY_TYPE_UC, 2) | \
                                  PAT_ENTRY(MEMORY_TYPE_UC, 3) | \
                                  PAT_ENTRY(MEMORY_TYPE_WB, 4) | \
                                  PAT_ENTRY(MEMORY_TYPE_WC, 5) | \
                                  PAT_ENTRY(MEMORY_TYPE_UC, 6) | \
                                  PAT_ENTRY(MEMORY_TYPE_UC, 7))

KCPU gCpus[KE_MAX_CPUS];
uint32_t gCpuCount = 0;

static uint8_t gBootstrapCpuStack[BOOTSTRAP_CPU_STACK_SIZE];
static volatile int32_t gKernelGsReady = 0;
static volatile int32_t gWriteCombiningAvailable = 0;

extern KCPU* KeReadCurrentCpuFromGs(void);

// Returns whether this CPU exposes the memory-type mechanisms required by the kernel policy.
static bool KiMemoryTypeFeaturesAvailable(void) {
    int regs[4];
    __cpuid(regs, 1);
    uint32_t features = (uint32_t)regs[3];
    return (features & (CPUID_FEATURE_MTRR | CPUID_FEATURE_PAT)) ==
           (CPUID_FEATURE_MTRR | CPUID_FEATURE_PAT);
}

// Flushes all translations before memory-type state changes and leaves global pages disabled.
static uintptr_t KiDisableGlobalTranslations(void) {
    uintptr_t cr4 = __readcr4();

    // nt-lite does not enable PCID. Refuse inherited PCID state rather than
    // pretending that a CR3 reload flushed every tagged translation.
    if (cr4 & CR4_PCID_ENABLE)
        KeStopProcessor();

    if (cr4 & CR4_PAGE_GLOBAL_ENABLE)
        __writecr4(cr4 & ~CR4_PAGE_GLOBAL_ENABLE);
    else
        __writecr3(__readcr3());

    return cr4;
}

// Returns whether this CPU accepted the complete canonical PAT and MTRR policy.
static bool KiMemoryTypePolicyInstalled(uint32_t variableCount) {
    if (__readmsr(IA32_PAT) != CANONICAL_PAT)
        return false;

    uint64_t expectedDefault = MTRR_DEF_ENABLE | MEMORY_TYPE_WB;
    uint64_t defaultType = __readmsr(IA32_MTRR_DEF_TYPE);
    uint64_t relevantDefaultBits = MTRR_DEF_TYPE_MASK | MTRR_DEF_FIXED_ENABLE | MTRR_DEF_ENABLE;
    if ((defaultType & relevantDefaultBits) != expectedDefault)
        return false;

    for (uint32_t i = 0; i < variableCount; i++) {
        uint64_t base = __readmsr(IA32_MTRR_PHYSBASE0 + i * 2u);
        uint64_t mask = __readmsr(IA32_MTRR_PHYSMASK0 + i * 2u);
        if (base != 0 || (mask & MTRR_MASK_VALID) != 0)
            return false;
    }

    return true;
}

// Programs the kernel-owned PAT policy over a uniform write-back MTRR baseline.
void KeInitializeMemoryTypes(void) {
    if (!KiMemoryTypeFeaturesAvailable())
        KeStopProcessor();

    uint32_t variableCount = (uint32_t)(__readmsr(IA32_MTRRCAP) & MTRR_CAP_VARIABLE_COUNT);
    if (KiMemoryTypePolicyInstalled(variableCount)) {
        NtlAtomicExchange32(&gWriteCombiningAvailable, 1);
        return;
    }

    uintptr_t interruptFlags = KeDisableInterrupts();
    uintptr_t cr0 = __readcr0();

    // Every CPU runs this before entering normal kernel execution. The boot processor
    // installs the policy while it is alone, and each AP installs it before
    // publishing its started state.
    __writecr0((cr0 | CR0_CACHE_DISABLE) & ~CR0_NOT_WRITE_THROUGH);
    __wbinvd();
    uintptr_t cr4 = KiDisableGlobalTranslations();

    // MTRRs provide only a uniform WB baseline. PAT is the sole owner of the
    // page-granular WB, WC, and UC policy selected by kernel mappings.
    __writemsr(IA32_MTRR_DEF_TYPE, 0);
    for (uint32_t i = 0; i < variableCount; i++) {
        __writemsr(IA32_MTRR_PHYSBASE0 + i * 2u, 0);
        __writemsr(IA32_MTRR_PHYSMASK0 + i * 2u, 0);
    }

    __writemsr(IA32_PAT, CANONICAL_PAT);
    __writemsr(IA32_MTRR_DEF_TYPE, MTRR_DEF_ENABLE | MEMORY_TYPE_WB);

    __wbinvd();
    __writecr3(__readcr3());
    __writecr0(cr0);
    if (cr4 & CR4_PAGE_GLOBAL_ENABLE)
        __writecr4(cr4);

    bool installed = KiMemoryTypePolicyInstalled(variableCount);
    KeRestoreInterrupts(interruptFlags);
    if (!installed)
        KeStopProcessor();

    NtlAtomicExchange32(&gWriteCombiningAvailable, 1);
}

// Reports whether PAT write-combining is available on the current platform.
bool KeWriteCombiningAvailable(void) {
    return NtlAtomicLoad32(&gWriteCombiningAvailable) != 0;
}

// Loads kernel Gs Base.
void KeLoadKernelGsBase(KCPU* cpu) {
    if (!cpu)
        cpu = &gCpus[0];

    cpu->Self = cpu;
    __writemsr(IA32_GS_BASE, (uint64_t)(uintptr_t)cpu);
    gKernelGsReady = 1;
}

// Updates user Gs Base.
void KeSetUserGsBase(uintptr_t userGsBase) {
    // While running in the kernel, GS_BASE is the current KCPU. The user TEB
    // base is parked in IA32_KERNEL_GS_BASE so swapgs can exchange it with the
    // KCPU pointer on return to ring 3.
    __writemsr(IA32_KERNEL_GS_BASE, (uint64_t)userGsBase);
}

// Returns user Gs Base.
uintptr_t KeGetUserGsBase(void) {
    return (uintptr_t)__readmsr(IA32_KERNEL_GS_BASE);
}

// Initializes bootstrap Cpu.
void KeInitializeBootstrapCpu(void) {
    memset(gCpus, 0, sizeof(gCpus));

    KCPU* cpu = &gCpus[0];
    cpu->Self = cpu;
    cpu->CpuIndex = 0;
    cpu->HardwareId = 0;
    cpu->State = KeCpuStarted;
    cpu->IsBootProcessor = 1;
    cpu->PackageId = 0;
    cpu->CoreId = 0;
    cpu->ThreadId = 0;
    cpu->PhysicalCoreIndex = 0;
    cpu->CoreSiblingMask = 1ULL;
    cpu->PackageMask = 1ULL;
    cpu->IsPrimaryThread = 1;
    cpu->SmtThreadCount = 1;
    cpu->TopologyFlags = 0;
    cpu->CurrentIrql = PASSIVE_LEVEL;
    cpu->InterruptDepth = 0;
    cpu->ReschedulePending = 0;
    cpu->KernelStack = gBootstrapCpuStack;
    cpu->KernelStackTop = ((uintptr_t)(gBootstrapCpuStack + sizeof(gBootstrapCpuStack))) & ~(uintptr_t)15;
    memset(&cpu->Tss, 0, sizeof(cpu->Tss));
    cpu->Tss.IoMapBase = sizeof(TSS64);
    cpu->Tss.RSP0 = cpu->KernelStackTop;

    gCpuCount = 1;
    KeInitializeMemoryTypes();
    KeLoadKernelGsBase(cpu);
    KeSetUserGsBase(0);
}

// Returns current Cpu.
KCPU* KeGetCurrentCpu(void) {
    if (gKernelGsReady) {
        KCPU* cpu = KeReadCurrentCpuFromGs();
        if (cpu >= &gCpus[0] && cpu < &gCpus[KE_MAX_CPUS] && cpu->Self == cpu)
            return cpu;
    }

    // Early bootstrap fallback before GS_BASE has been installed.
    return &gCpus[0];
}

// Returns cpu By Index.
KCPU* KeGetCpuByIndex(uint32_t index) {
    if (index >= KE_MAX_CPUS)
        return (KCPU*)0;
    return &gCpus[index];
}

// Returns current Cpu Index.
uint32_t KeGetCurrentCpuIndex(void) {
    return KeGetCurrentCpu()->CpuIndex;
}

TSS64* KeGetCurrentTss(void) {
    return &KeGetCurrentCpu()->Tss;
}