// Kernel statistics ABI smoke test.
// Verifies size queries, trailing CPU records, and generic filesystem stats.

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <ntlite/syscall.h>

// Prints one NUL-terminated string to the kernel console.
static void print(const char* s) {
    SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromCstr(s) };
    SysWriteConsole(&req);
}

// Prints one unsigned integer to the kernel console.
static void print_u64(uint64_t v) {
    char buf[32];
    int i = 0;
    if (v == 0) {
        print("0");
        return;
    }
    while (v && i < (int)sizeof(buf)) {
        buf[i++] = (char)('0' + (v % 10));
        v /= 10;
    }
    while (i--) {
        char c = buf[i];
        SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromParts(&c, 1) };
        SysWriteConsole(&req);
    }
}

// Terminates the current test process.
static void exit_now(int code) {
    while (true) SysExitProcess(code);
}

// Fails the test with one plain message.
static void fail(const char* what) {
    print("stats_abi_smoke: ");
    print(what);
    print("\n");
    exit_now(1);
}

// Fails the test with one status value.
static void fail_status(const char* what, uint32_t status) {
    print("stats_abi_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Requires one exact status value.
static void expect_status(const char* what, uint32_t got, uint32_t expected) {
    if (got == expected) return;
    print("stats_abi_smoke: ");
    print(what);
    print(" got=");
    print_u64(got);
    print(" expected=");
    print_u64(expected);
    print("\n");
    exit_now(1);
}

// Verifies fixed-size stats classes report and fill their public header.
static void check_fixed_stats(void) {
    size_t required = 0;
    uint32_t status = SysQueryKernelStats(SYS_KERNEL_STATS_SCHEDULER, 0, 0, &required);
    if (status != NTLSTATUS_SUCCESS) fail_status("scheduler size", status);
    if (required != sizeof(SYS_SCHEDULER_STATS)) fail("scheduler required size mismatch");

    uint8_t small[sizeof(SYS_SCHEDULER_STATS) - 1u];
    size_t smallRequired = 0;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_SCHEDULER, small, sizeof(small), &smallRequired);
    expect_status("scheduler small buffer", status, NTLSTATUS_BUFFER_TOO_SMALL);
    if (smallRequired != sizeof(SYS_SCHEDULER_STATS)) fail("scheduler small required mismatch");

    SYS_SCHEDULER_STATS scheduler;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_SCHEDULER, &scheduler, sizeof(scheduler), 0);
    if (status != NTLSTATUS_SUCCESS) fail_status("scheduler exact", status);
    if (scheduler.Size != sizeof(scheduler)) fail("scheduler size field mismatch");
    if (scheduler.Version != SYS_STATS_VERSION) fail("scheduler version mismatch");

    SYS_FILESYSTEM_STATS filesystem;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_FILESYSTEM, &filesystem, sizeof(filesystem), 0);
    if (status != NTLSTATUS_SUCCESS) fail_status("filesystem exact", status);
    if (filesystem.Size != sizeof(filesystem)) fail("filesystem size field mismatch");
    if (filesystem.Version != SYS_STATS_VERSION) fail("filesystem version mismatch");

    SYS_EXCEPTION_STATS exceptionStats;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_EXCEPTION, &exceptionStats, sizeof(exceptionStats), 0);
    if (status != NTLSTATUS_SUCCESS) fail_status("exception exact", status);
    if (exceptionStats.Size != sizeof(exceptionStats)) fail("exception size field mismatch");
    if (exceptionStats.Version != SYS_STATS_VERSION) fail("exception version mismatch");
}

// Verifies CPU stats use a size-reported trailing-record layout.
static void check_cpu_stats(void) {
    size_t required = 0;
    uint32_t status = SysQueryKernelStats(SYS_KERNEL_STATS_CPU, 0, 0, &required);
    if (status != NTLSTATUS_SUCCESS) fail_status("cpu size", status);
    if (required < sizeof(SYS_CPU_STATS) + sizeof(SYS_CPU_STAT)) fail("cpu required size too small");

    SYS_CPU_STATS headerOnly;
    size_t smallRequired = 0;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_CPU, &headerOnly, sizeof(headerOnly), &smallRequired);
    expect_status("cpu header-only buffer", status, NTLSTATUS_BUFFER_TOO_SMALL);
    if (smallRequired != required) fail("cpu small required mismatch");

    void* buffer = 0;
    size_t pages = (required + 4095u) / 4096u;
    if (pages == 0) pages = 1;
    SYS_ALLOC_VM_REQUEST allocReq;
    memset(&allocReq, 0, sizeof(allocReq));
    allocReq.Size = sizeof(allocReq);
    allocReq.PageCount = pages;
    allocReq.Protect = SYS_VM_READWRITE;
    allocReq.Flags = SYS_VM_FLAG_NONE;
    allocReq.BaseOut = &buffer;
    status = SysAllocVm(&allocReq);
    if (status != NTLSTATUS_SUCCESS || !buffer) fail_status("cpu buffer alloc", status);

    SYS_CPU_STATS* cpu = (SYS_CPU_STATS*)buffer;
    status = SysQueryKernelStats(SYS_KERNEL_STATS_CPU, cpu, pages * 4096u, 0);
    if (status != NTLSTATUS_SUCCESS) {
        SysFreeVm(buffer);
        fail_status("cpu exact", status);
    }

    if (cpu->Size != required) fail("cpu size field mismatch");
    if (cpu->Version != SYS_STATS_VERSION) fail("cpu version mismatch");
    if (cpu->HeaderSize != sizeof(SYS_CPU_STATS)) fail("cpu header size mismatch");
    if (cpu->RecordSize != sizeof(SYS_CPU_STAT)) fail("cpu record size mismatch");
    if (cpu->RecordCount == 0 || cpu->RecordCount > cpu->CpuCount) fail("cpu record count mismatch");

    const SYS_CPU_STAT* records = SYS_CPU_STATS_CONST_RECORDS(cpu);
    if (records[0].CpuIndex >= cpu->CpuCount) fail("cpu first record index mismatch");

    SysFreeVm(buffer);
}

// Runs the stats ABI smoke test.
int main(void) {
    print("stats_abi_smoke: start\n");
    check_fixed_stats();
    check_cpu_stats();
    print("stats_abi_smoke: PASS\n");
    exit_now(0);
    return 0;
}