// Focused debug-only checks for bootstrap state, APC waits, and copy-on-write.

#include <stdint.h>
#include <intrin.h>
#include <ntlite/status.h>
#include <kernel/ke_apc.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_event.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_process.h>
#include <kernel/ob_object.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pfn.h>

static void KiAssertApcNormalContext(const char* name) {
    KCPU* cpu = KeGetCurrentCpu();
    if (KeGetCurrentIrql() != PASSIVE_LEVEL)
        KePanic(KE_PANIC_CODE_SELFTEST, "%s ran above PASSIVE_LEVEL\n", name);
    if ((__readeflags() & 0x200u) == 0)
        KePanic(KE_PANIC_CODE_SELFTEST, "%s ran with interrupts disabled\n", name);
    if (!cpu || cpu->InterruptDepth != 0)
        KePanic(KE_PANIC_CODE_SELFTEST, "%s ran inside interrupt context depth=%u\n",
                name, cpu ? cpu->InterruptDepth : 0xFFFFFFFFu);
}

// Leaves the APC's normal routine and arguments unchanged.
static void KiApcWaitRaceKernelRoutine(KAPC* apc,
                                       void** normalRoutine,
                                       void** normalContext,
                                       void** systemArgument1,
                                       void** systemArgument2) {
    (void)apc;
    (void)normalRoutine;
    (void)normalContext;
    (void)systemArgument1;
    (void)systemArgument2;
}

// Records delivery for the prepared-wait race probe.
static void KiApcWaitRaceNormalRoutine(void* normalContext,
                                       void* systemArgument1,
                                       void* systemArgument2) {
    (void)systemArgument1;
    (void)systemArgument2;
    KiAssertApcNormalContext("APC wait-race normal routine");
    (*(volatile int*)normalContext)++;
}

// Reproduces the prepare-to-commit APC race deterministically.
static void KiRunApcWaitRaceSelfTest(void) {
    DbgPrint("APC wait-race selftest: begin\n");

    THREAD* self = KeGetCurrentThread();
    if (!self || KeGetCurrentIrql() != PASSIVE_LEVEL)
        KePanic(KE_PANIC_CODE_SELFTEST, "APC wait-race selftest: invalid execution context\n");

    EVENT* waitEvent = (EVENT*)0;
    KAPC apc;
    volatile int delivered = 0;
    KSTATUS status = KeEventCreate(EventAutoReset, 0, &waitEvent);
    if (!KSUCCESS(status))
        KePanic(KE_PANIC_CODE_SELFTEST, "APC wait-race selftest: event create failed %u\n", status);

    KeInitializeApc(&apc, self, KiApcWaitRaceKernelRoutine,
                    KiApcWaitRaceNormalRoutine, (void*)&delivered);

    uintptr_t flags = KeAcquireSpinlock(&waitEvent->Lock);
    if (!KeInsertQueueApc(&apc, (void*)(uintptr_t)0x99999999u,
                          (void*)(uintptr_t)0xAAAAAAAAu))
        KePanic(KE_PANIC_CODE_SELFTEST, "APC wait-race selftest: queue failed\n");
    if (delivered != 0)
        KePanic(KE_PANIC_CODE_SELFTEST, "APC wait-race selftest: delivered above PASSIVE_LEVEL\n");

    KePrepareBlockThread(self, &waitEvent->Header, 1);
    KeListInsertTailChecked(&waitEvent->WaitList, &self->SchedulerLink);
    KeSetThreadWaitList(self, &waitEvent->WaitList);
    KeReleaseSpinlock(&waitEvent->Lock, flags);

    status = KeCommitBlockThread(self, 1);
    if (status != NTLSTATUS_KERNEL_APC || delivered != 1)
        KePanic(KE_PANIC_CODE_SELFTEST,
                "APC wait-race selftest: bad result status=%u delivered=%d\n",
                status, delivered);

    ObDereferenceObject(&waitEvent->Header);
    DbgPrint("APC wait-race selftest: PASS\n");
}

// Verifies that writable copy-on-write views split before either alias can modify shared state.
static void KiRunVadCopyOnWriteSelfTest(void) {
    PROCESS* process = KeGetCurrentProcess();
    if (!process)
        process = KeGetSystemProcess();
    if (!process)
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: no process context\n");

    void* base = (void*)0;
    KSTATUS s = MmVadCreateCommitted(process,
                                     PAGE_SIZE,
                                     PAGE_SIZE,
                                     PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                                     MM_VAD_ZERO,
                                     MmVadGeneric,
                                     &base);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: source allocation failed %u\n", s);

    s = MmVadEnsureResident(process, base, PAGE_SIZE, PAGE_WRITE | PAGE_USER);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: source materialization failed %u\n", s);

    uint8_t* source = (uint8_t*)base;
    source[0] = 0x31;

    void* cowBase = (void*)0;
    s = MmVadCreateCowView(process, base, PAGE_SIZE, process, &cowBase);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: COW view creation failed %u\n", s);

    PHYSICAL_ADDRESS sourcePhys = 0;
    PHYSICAL_ADDRESS cowPhys = 0;
    uint64_t sourceFlags = 0;
    uint64_t cowFlags = 0;
    s = MmPteQueryPageFor(MmPteQuerySelfTest, (uintptr_t)base, &sourcePhys, &sourceFlags);
    if (!KSUCCESS(s) || (sourceFlags & PAGE_WRITE))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: source was not write-protected status=%u flags=%016llx\n",
                s, (unsigned long long)sourceFlags);

    s = MmPteQueryPageFor(MmPteQuerySelfTest, (uintptr_t)cowBase, &cowPhys, &cowFlags);
    if (!KSUCCESS(s) || (cowFlags & PAGE_WRITE) || cowPhys != sourcePhys)
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: view was not shared read-only status=%u src=%016llx cow=%016llx flags=%016llx\n",
                s,
                (unsigned long long)sourcePhys,
                (unsigned long long)cowPhys,
                (unsigned long long)cowFlags);

    s = MmVadResolveUserPageFault(process, (uintptr_t)cowBase, 0x3);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: simulated write fault split failed %u\n", s);

    uint8_t* copy = (uint8_t*)cowBase;
    copy[0] = 0x62;
    if (source[0] != 0x31 || copy[0] != 0x62)
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: copy isolation failed src=%02x cow=%02x\n",
                source[0], copy[0]);

    PHYSICAL_ADDRESS splitPhys = 0;
    uint64_t splitFlags = 0;
    s = MmPteQueryPageFor(MmPteQuerySelfTest, (uintptr_t)cowBase, &splitPhys, &splitFlags);
    if (!KSUCCESS(s) || splitPhys == sourcePhys || !(splitFlags & PAGE_WRITE))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: split mapping wrong status=%u old=%016llx new=%016llx flags=%016llx\n",
                s,
                (unsigned long long)sourcePhys,
                (unsigned long long)splitPhys,
                (unsigned long long)splitFlags);

    (void)MmVadDestroy(process, cowBase);
    (void)MmVadDestroy(process, base);

    void* zeroBase = (void*)0;
    s = MmVadCreateCommitted(process,
                             PAGE_SIZE,
                             PAGE_SIZE,
                             PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                             MM_VAD_ZERO,
                             MmVadGeneric,
                             &zeroBase);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero-page VAD creation failed %u\n", s);

    s = MmVadEnsureResident(process, zeroBase, PAGE_SIZE, PAGE_USER);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero-page read materialization failed %u\n", s);

    PHYSICAL_ADDRESS zeroPhys = 0;
    uint64_t zeroFlags = 0;
    s = MmPteQueryPageFor(MmPteQuerySelfTest, (uintptr_t)zeroBase, &zeroPhys, &zeroFlags);
    if (!KSUCCESS(s) || (zeroFlags & PAGE_WRITE))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero page was not mapped read-only status=%u flags=%016llx\n",
                s, (unsigned long long)zeroFlags);

    s = MmVadResolveUserPageFault(process, (uintptr_t)zeroBase, 0x3);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero-page split failed %u\n", s);

    PHYSICAL_ADDRESS privatePhys = 0;
    uint64_t privateFlags = 0;
    s = MmPteQueryPageFor(MmPteQuerySelfTest, (uintptr_t)zeroBase, &privatePhys, &privateFlags);
    if (!KSUCCESS(s) || privatePhys == zeroPhys || !(privateFlags & PAGE_WRITE))
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero-page split mapping wrong status=%u old=%016llx new=%016llx flags=%016llx\n",
                s,
                (unsigned long long)zeroPhys,
                (unsigned long long)privatePhys,
                (unsigned long long)privateFlags);

    ((uint8_t*)zeroBase)[0] = 0x9C;
    if (((uint8_t*)zeroBase)[0] != 0x9C)
        KePanic(KE_PANIC_CODE_SELFTEST, "VAD COW selftest: zero-page private write failed\n");

    (void)MmVadDestroy(process, zeroBase);
    DbgPrint("VAD COW selftest: shared views and zero-page split PASS\n");
}

void KiCheckBootState(void) {
    KCPU* cpu = KeGetCurrentCpu();
    THREAD* thread = KeGetCurrentThread();
    if (!cpu || gCpuCount == 0 || gCpuCount > KE_MAX_CPUS ||
        cpu != &gCpus[0] || cpu->Self != cpu ||
        cpu->State != KeCpuOnline || !cpu->IsBootProcessor ||
        KeGetCurrentIrql() != PASSIVE_LEVEL ||
        !thread || thread != cpu->CurrentThread ||
        thread != cpu->IdleThread || !thread->IsIdleThread)
        KePanic(KE_PANIC_CODE_SELFTEST, "boot state: invalid bootstrap CPU or idle thread\n");
}

void KiRunKernelSelfTests(void) {
    DbgPrint("Kernel diagnostics: begin\n");
    KiRunApcWaitRaceSelfTest();
    KiRunVadCopyOnWriteSelfTest();
    DbgPrint("Kernel diagnostics: PASS\n");
}