// Process-construction smoke test. The same import-free PE is loaded a second time entirely by userspace.

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

#define TEST_STARTUP_EVENT_TAG 0x50425445u

static void exit_now(int code) {
    for (;;) SysExitProcess(code);
}

// Writes one string through the native console syscall.
static void print(const char* text) {
    SYS_WRITE_CONSOLE_REQUEST request = { sizeof(request), SysSvFromCstr(text) };
    (void)SysWriteConsole(&request);
}

// Writes one 32-bit value as hexadecimal without requiring a formatting runtime.
static void print_hex32(uint32_t value) {
    static const char digits[] = "0123456789ABCDEF";
    char buffer[11] = "0x00000000";
    for (uint32_t i = 0; i < 8; ++i)
        buffer[2 + i] = digits[(value >> ((7u - i) * 4u)) & 0xFu];
    print(buffer);
}

static int is_child(void) {
    SYS_PEB* peb = SysGetPeb();
    static const char marker[] = "process-builder-child";
    return peb && peb->CommandLine.Length == sizeof(marker) - 1 &&
           memcmp(peb->CommandLine.Data, marker, sizeof(marker) - 1) == 0;
}

static SYS_HANDLE find_startup(uint32_t tag) {
    SYS_PEB* peb = SysGetPeb();
    if (!peb || peb->Magic != SYS_PEB_MAGIC || peb->Version != SYS_ENV_VERSION)
        return SYS_HANDLE_INVALID;
    for (uint32_t i = 0; i < peb->StartupCapabilityCount; ++i) {
        if (peb->StartupCapabilities[i].Tag == tag)
            return peb->StartupCapabilities[i].Handle;
    }
    return SYS_HANDLE_INVALID;
}

int main(void) {
    if (is_child()) {
        SYS_HANDLE event = find_startup(TEST_STARTUP_EVENT_TAG);
        if (event == SYS_HANDLE_INVALID) exit_now(10);
        if (SysSignalEvent(event) != NTLSTATUS_SUCCESS) exit_now(11);
        exit_now(0);
    }

    SYS_PEB* peb = SysGetPeb();
    if (!peb || peb->Magic != SYS_PEB_MAGIC || peb->Version != SYS_ENV_VERSION)
        exit_now(1);

    SYS_HANDLE file = SYS_HANDLE_INVALID;
    KSTATUS s = NtlOpenFile(peb->ImagePath,
                                                        SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY,
                                                        SYS_FILE_OPEN_EXISTING,
                                                        SYS_FILE_FLAG_NONE,
                                         &file);
    if (!KSUCCESS(s)) exit_now(2);

    SYS_HANDLE child = SYS_HANDLE_INVALID;
    SYS_CREATE_PROCESS_OBJECT_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Flags = SYS_CREATE_PROCESS_OBJECT_FLAG_NONE;
    create.ImagePath = peb->ImagePath;
    create.CommandLine = SYS_SV_LITERAL("process-builder-child");
    create.SecurityContext = SYS_SECURITY_CONTEXT_CURRENT;
    create.ProcessHandleOut = &child;
    s = SysCreateProcessObject(&create);
    if (!KSUCCESS(s)) exit_now(3);

    NTL_PE_IMAGE image;
    memset(&image, 0, sizeof(image));
    NTL_PE_LOAD_REQUEST load;
    memset(&load, 0, sizeof(load));
    load.Size = sizeof(load);
    load.ProcessHandle = child;
    load.FileHandle = file;
    load.Flags = NTL_PE_LOAD_FLAG_NONE;
    load.ImageOut = &image;
    s = NtlLoadPeImage(&load);
    SysCloseHandle(file);
    if (!KSUCCESS(s) || !image.Base || !image.EntryPoint) {
        print("process_builder_smoke: NtlLoadPeImage failed status=");
        print_hex32(s);
        print("\n");
        exit_now(4);
    }

    SYS_HANDLE event = SYS_HANDLE_INVALID;
    s = SysCreateEvent(0, 0, &event);
    if (!KSUCCESS(s)) exit_now(5);

    SYS_HANDLE childEvent = SYS_HANDLE_INVALID;
    SYS_DUPLICATE_HANDLE_REQUEST duplicate;
    memset(&duplicate, 0, sizeof(duplicate));
    duplicate.Size = sizeof(duplicate);
    duplicate.SourceHandle = event;
    duplicate.TargetProcessHandle = child;
    duplicate.DesiredRights = SYS_EVENT_RIGHT_SIGNAL | SYS_RIGHT_QUERY;
    duplicate.Flags = SYS_DUPLICATE_HANDLE_FLAG_NONE;
    duplicate.TargetHandleOut = &childEvent;
    s = SysDuplicateHandle(&duplicate);
    if (!KSUCCESS(s)) exit_now(6);

    SYS_STARTUP_CAPABILITY startup = { TEST_STARTUP_EVENT_TAG, childEvent };
    SYS_START_PROCESS_REQUEST start;
    memset(&start, 0, sizeof(start));
    start.Size = sizeof(start);
    start.ProcessHandle = child;
    start.Flags = SYS_START_PROCESS_FLAG_NONE;
    start.MainImageBase = image.Base;
    start.MainImageSize = image.Size;
    start.EntryPoint = image.EntryPoint;
    start.StartupCapabilities = &startup;
    start.StartupCapabilityCount = 1;
    s = SysStartProcess(&start);
    if (!KSUCCESS(s)) exit_now(7);

    s = SysWait(event, SYS_WAIT_FLAG_NONE);
    if (!KSUCCESS(s)) exit_now(8);
    s = SysWait(child, SYS_WAIT_FLAG_NONE);
    if (!KSUCCESS(s)) exit_now(9);
    int32_t exitCode = -1;
    s = SysGetProcessExitCode(child, &exitCode);
    SysCloseHandle(event);
    SysCloseHandle(child);
    if (!KSUCCESS(s) || exitCode != 0) exit_now(12);
    return 0;
}