// ke_cpu.c - bootstrap/per-CPU state plumbing
#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);

void KeLoadKernelGsBase(KCPU* cpu) {
    if (!cpu)
        cpu = &gCpus[0];

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

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

uintptr_t KeGetUserGsBase(void) {
    return (uintptr_t)__readmsr(IA32_KERNEL_GS_BASE);
}

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

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];
}

KCPU* KeGetCpuByIndex(uint32_t index) {
    if (index >= KE_MAX_CPUS)
        return (KCPU*)0;
    return &gCpus[index];
}

uint32_t KeGetCurrentCpuIndex(void) {
    return KeGetCurrentCpu()->CpuIndex;
}

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