// User-mode smoke test for process-local userspace module loading.

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

static void print(const char* s) {
    SYS_WRITE_CONSOLE_REQUEST req = { sizeof(SYS_WRITE_CONSOLE_REQUEST), SysSvFromCstr(s) };
    SysWriteConsole(&req);
}

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 = { sizeof(req), SysSvFromParts(&c, 1) };
        SysWriteConsole(&req);
    }
}

static void exit_now(int32_t code) { while (true) SysExitProcess(code); }
static void fail(const char* message) { print("module_load_smoke: "); print(message); print("\n"); exit_now(1); }
static void fail_status(const char* what, KSTATUS status) {
    print("module_load_smoke: "); print(what); print(" status="); print_u64(status); print("\n"); exit_now(1);
}

static void check_module_info(const NTL_MODULE_INFO* info) {
    if (!info || !info->Base || !info->Size) fail("module has empty mapping");
    if (info->PublicRefCount == 0) fail("module public reference missing");
    if (strcmp(info->Name, "posix.dll") != 0) fail("module name mismatch");
}

static bool module_list_contains(uintptr_t base) {
    size_t count = 0;
    KSTATUS status = NtlGetModuleList(0, 0, &count);
    if (status != NTLSTATUS_BUFFER_TOO_SMALL && status != NTLSTATUS_SUCCESS)
        fail_status("query module count", status);
    if (count == 0) return false;
    if (count > SIZE_MAX / sizeof(NTL_MODULE_INFO)) fail("module count overflow");
    size_t bytes = count * sizeof(NTL_MODULE_INFO);
    size_t pages = (bytes + 4095u) / 4096u;
    void* memory = 0;
    SYS_ALLOC_VM_REQUEST alloc;
    memset(&alloc, 0, sizeof(alloc));
    alloc.Size = sizeof(alloc);
    alloc.Protect = SYS_VM_READWRITE;
    alloc.Flags = SYS_VM_FLAG_NONE;
    alloc.PageCount = pages;
    alloc.BaseOut = &memory;
    status = SysAllocVm(&alloc);
    if (status != NTLSTATUS_SUCCESS) fail_status("allocate module list", status);
    NTL_MODULE_INFO* modules = (NTL_MODULE_INFO*)memory;
    memset(modules, 0, bytes);
    size_t copied = count;
    status = NtlGetModuleList(modules, count, &copied);
    if (status != NTLSTATUS_SUCCESS) { (void)SysFreeVm(memory); fail_status("query module list", status); }
    bool found = false;
    for (size_t i = 0; i < copied; ++i) if (modules[i].Base == base) { found = true; break; }
    (void)SysFreeVm(memory);
    return found;
}

// Loads one module through the userspace registry with the specified filename policy.
static NTL_MODULE_HANDLE load_module(SYS_HANDLE root, const char* path,
                                     NTL_MODULE_NAME_RESOLVER resolver, NTL_MODULE_INFO* info) {
    NTL_MODULE_HANDLE handle = NTL_MODULE_HANDLE_INVALID;
    NTL_LOAD_MODULE_AT_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.Flags = NTL_MODULE_LOAD_FLAG_NONE;
    req.RootHandle = root;
    req.Path = SysSvFromCstr(path);
    req.ModuleRootHandle = root;
    req.ModuleDirectory = SYS_SV_LITERAL("system");
    req.ResolveModuleName = resolver;
    req.ModuleHandleOut = &handle;
    req.InfoOut = info;
    KSTATUS status = NtlLoadModuleAt(&req);
    if (status != NTLSTATUS_SUCCESS) fail_status("load module", status);
    if (handle == NTL_MODULE_HANDLE_INVALID) fail("load returned null handle");
    return handle;
}

// Runs the userspace process-module loader smoke test.
int main(void) {
    print("module_load_smoke: start\n");
    SYS_HANDLE root = SYS_HANDLE_INVALID;
    KSTATUS status = NtlResourceOpenSystemRoot(SYS_RIGHT_QUERY | SYS_VOLUME_RIGHT_OPEN, &root);
    if (status != NTLSTATUS_SUCCESS) fail_status("open system root", status);

    NTL_MODULE_HANDLE syscallModule = load_module(root, "system/syscall.dll",
                                                  NtlResolveModuleNameIdentity, 0);
    uintptr_t address = 0;
    status = NtlGetModuleExport(syscallModule, SysSvFromCstr("SysGetCurrentPid"), &address);
    if (status != NTLSTATUS_SUCCESS || !address) fail_status("lookup SysGetCurrentPid", status);
    status = NtlUnloadModule(syscallModule);
    if (status != NTLSTATUS_SUCCESS) fail_status("unload syscall.dll", status);

    NTL_MODULE_INFO info;
    memset(&info, 0, sizeof(info));
    NTL_MODULE_HANDLE posix = load_module(root, "system/msvcrt.dll",
                                          NtlResolveSystemModuleName, &info);
    check_module_info(&info);
    if (!module_list_contains(info.Base)) fail("loaded module absent from list");

    NTL_MODULE_INFO byHandle;
    memset(&byHandle, 0, sizeof(byHandle));
    status = NtlGetModuleInfo(posix, &byHandle);
    if (status != NTLSTATUS_SUCCESS) fail_status("get module info", status);
    check_module_info(&byHandle);
    if (byHandle.Base != info.Base) fail("module handle base mismatch");

    status = NtlUnloadModule(posix);
    if (status != NTLSTATUS_SUCCESS) fail_status("unload posix.dll", status);
    status = NtlGetModuleInfo(posix, &byHandle);
    if (status != NTLSTATUS_INVALID_HANDLE) fail_status("query unloaded handle", status);

    (void)SysCloseHandle(root);
    print("module_load_smoke: PASS\n");
    exit_now(0);
    return 0;
}