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

#include <stdint.h>
#include <string.h>
#include <intrin.h>
#include <kernel/ke_cpu.h>

#define BOOTSTRAP_CPU_STACK_SIZE 0x4000u
#define IA32_GS_BASE             0xC0000101u
#define IA32_KERNEL_GS_BASE      0xC0000102u

KCPU gCpus[KE_MAX_CPUS];
uint32_t gCpuCount = 0;

static uint8_t gBootstrapCpuStack[BOOTSTRAP_CPU_STACK_SIZE];
static volatile long gKernelGsReady = 0;

extern KCPU* KeReadCurrentCpuFromGs(void);

// 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->LapicId = 0;
    cpu->Present = 1;
    cpu->Started = 1;
    cpu->Parked = 0;
    cpu->Online = 1;
    cpu->IsBsp = 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->SchedulerCommand = 0;
    cpu->SchedulerActive = 0;
    cpu->KernelStack = gBootstrapCpuStack;
    cpu->KernelStackTop = (uintptr_t)(gBootstrapCpuStack + sizeof(gBootstrapCpuStack));
    memset(&cpu->Tss, 0, sizeof(cpu->Tss));
    cpu->Tss.IoMapBase = sizeof(TSS64);
    cpu->Tss.RSP0 = cpu->KernelStackTop;

    gCpuCount = 1;
    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;
}