// Native userspace process construction built on capability-relative file open and PE materialization.

#include <stdbool.h>
#include <ntlite/process_loader.h>
#include <ntlite/path.h>
#include <ntlite/resource.h>
#include <ntlite/syscall.h>
#include <string.h>
#include "module_loader_internal.h"

// Terminates one failed constructing child and releases the parent process handle.
static void NtlDiscardChild(SYS_HANDLE processHandle) {
    if (processHandle == SYS_HANDLE_INVALID)
        return;
    (void)SysKillProcess(processHandle, (int32_t)NTLSTATUS_PROCESS_KILLED);
    (void)SysCloseHandle(processHandle);
}

#define NTL_PREPARATION_PAGE_SIZE 4096u

// The preparation is allocated by the loader; its private layout is not ABI.
struct NTL_PROCESS_PREPARATION {
    uint32_t StartupCapabilityCount;
    SYS_HANDLE ProcessHandle;
    NTL_MODULE_CONTEXT* Modules;
    NTL_PE_IMAGE Image;
    void* LoaderData;
    SYS_STARTUP_CAPABILITY ChildCapabilities[SYS_STARTUP_CAPABILITY_MAX];
};

// Allocates zero-initialized page-backed preparation state in user space.
static KSTATUS NtlAllocatePreparation(NTL_PROCESS_PREPARATION** preparationOut) {
    if (!preparationOut) return NTLSTATUS_INVALID_PARAMETER;
    *preparationOut = 0;
    if (sizeof(NTL_PROCESS_PREPARATION) > SIZE_MAX - (NTL_PREPARATION_PAGE_SIZE - 1u))
        return NTLSTATUS_NO_MEMORY;
    size_t pages = (sizeof(NTL_PROCESS_PREPARATION) + NTL_PREPARATION_PAGE_SIZE - 1u) /
                   NTL_PREPARATION_PAGE_SIZE;
    SYS_ALLOC_VM_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Protect = SYS_VM_READWRITE;
    request.Flags = SYS_VM_FLAG_NONE;
    request.PageCount = pages;
    request.BaseOut = (void**)preparationOut;
    KSTATUS status = SysAllocVm(&request);
    if (KSUCCESS(status)) {
        memset(*preparationOut, 0, pages * NTL_PREPARATION_PAGE_SIZE);
        (*preparationOut)->ProcessHandle = SYS_HANDLE_INVALID;
    }
    return status;
}

static bool NtlPreparationRequestValid(const NTL_PREPARE_PROCESS_REQUEST* request) {
    return request && request->Size == sizeof(*request) &&
           request->Flags == NTL_PROCESS_LOAD_FLAG_NONE &&
           request->RootHandle != SYS_HANDLE_INVALID &&
           request->ModuleRootHandle != SYS_HANDLE_INVALID &&
           request->Path.Data && request->Path.Length != 0 &&
           SysSvIsValid(request->ModuleDirectory) &&
           request->StartupCapabilityCount <= SYS_STARTUP_CAPABILITY_MAX &&
           (!request->StartupCapabilityCount || request->StartupCapabilities) &&
           (!request->Arguments.Count || request->Arguments.Entries) &&
           (!request->Environment.Count || request->Environment.Entries) &&
           request->StartupDataSize <= NTL_PROCESS_STARTUP_DATA_MAX &&
           (!request->StartupDataSize || request->StartupData);
}

// Disposes private preparation state and optionally discards the inert child.
static void NtlDestroyPreparation(NTL_PROCESS_PREPARATION* preparation, bool discardChild) {
    if (!preparation) return;
    if (preparation->Modules) NtlpDestroyModuleContext(preparation->Modules);
    if (discardChild && preparation->ProcessHandle != SYS_HANDLE_INVALID)
        NtlDiscardChild(preparation->ProcessHandle);
    (void)SysFreeVm(preparation);
}

// Aborts and consumes an inert child; NULL requires no action.
void NtlAbortPreparedProcess(NTL_PROCESS_PREPARATION* preparation) {
    NtlDestroyPreparation(preparation, true);
}

// Opens the executable's containing directory as the private dependency-search root.
static KSTATUS NtlOpenImageDirectory(const NTL_PREPARE_PROCESS_REQUEST* request,
                                     SYS_HANDLE* directoryOut, int* borrowedRootOut) {
    if (!request || !directoryOut || !borrowedRootOut || !SysSvIsValid(request->Path) ||
        request->Path.Length == 0 || SysSvContainsByte(request->Path, '\0'))
        return NTLSTATUS_INVALID_PARAMETER;
    *directoryOut = SYS_HANDLE_INVALID;
    *borrowedRootOut = 0;

    size_t slash = SIZE_MAX;
    for (size_t i = 0; i < request->Path.Length; ++i) {
        if (request->Path.Data[i] == '/') slash = i;
        else if (request->Path.Data[i] == '\\') return NTLSTATUS_INVALID_PARAMETER;
    }
    if (slash == SIZE_MAX) {
        *directoryOut = request->RootHandle;
        *borrowedRootOut = 1;
        return NTLSTATUS_SUCCESS;
    }

    SYS_STRING_VIEW directoryPath;
    if (slash == 0) directoryPath = SYS_SV_LITERAL("/");
    else directoryPath = SysSvFromParts(request->Path.Data, slash);

    if (request->Path.Data[0] == '/')
        return NtlOpenDir(directoryPath, SYS_DIR_ACCESS_OPEN, SYS_DIR_FLAG_NONE, directoryOut);

    SYS_OPEN_DIR_AT_REQUEST open = SysMakeOpenDirAtRequest(request->RootHandle, directoryPath,
                                                           SYS_DIR_ACCESS_OPEN, SYS_DIR_FLAG_NONE,
                                                           directoryOut);
    return SysOpenDirAt(&open);
}

// Prepares an inert child: process object, image/dependencies, inherited startup
// capabilities, CWD, and loader/bootstrap block are all complete before return.
KSTATUS NtlPrepareProcess(const NTL_PREPARE_PROCESS_REQUEST* request,
                          SYS_HANDLE imageFile,
                          NTL_PROCESS_PREPARATION** preparationOut) {
    if (!preparationOut) return NTLSTATUS_INVALID_PARAMETER;
    *preparationOut = 0;
    if (!NtlPreparationRequestValid(request)) return NTLSTATUS_INVALID_PARAMETER;
    NTL_PROCESS_PREPARATION* state = 0;
    KSTATUS status = NtlAllocatePreparation(&state);
    if (!KSUCCESS(status)) return status;

    SYS_HANDLE openedImage = SYS_HANDLE_INVALID;
    if (imageFile == SYS_HANDLE_INVALID) {
        SYS_OPEN_FILE_AT_REQUEST open;
        memset(&open, 0, sizeof(open));
        open.Size = sizeof(open);
        open.RootHandle = request->RootHandle;
        open.Path = request->Path;
        open.Access = SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY;
        open.Disposition = SYS_FILE_OPEN_EXISTING;
        open.Flags = SYS_FILE_FLAG_NONE;
        open.HandleOut = &openedImage;
        KSTATUS openStatus = SysOpenFileAt(&open);
        if (!KSUCCESS(openStatus)) {
            NtlDestroyPreparation(state, true);
            return openStatus;
        }
        imageFile = openedImage;
    }

    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 = request->Path;
    create.CommandLine = request->CommandLine;
    create.SecurityContext = request->SecurityContext;
    create.ProcessHandleOut = &state->ProcessHandle;
    status = SysCreateProcessObject(&create);
    if (!KSUCCESS(status)) goto fail;

    char currentDirectoryBuffer[SYS_PATH_MAX + 1u];
    SYS_STRING_VIEW currentDirectory;
    size_t currentDirectorySize = 0;
    if (request->CurrentDirectory.Length != 0) {
        status = NtlRealPath(request->CurrentDirectory, currentDirectoryBuffer,
                             sizeof(currentDirectoryBuffer), &currentDirectorySize);
    } else {
        status = NtlGetCurrentDirectory(currentDirectoryBuffer, sizeof(currentDirectoryBuffer),
                                        &currentDirectorySize);
    }
    if (!KSUCCESS(status)) goto fail;
    currentDirectory = SysSvFromParts(currentDirectoryBuffer, currentDirectorySize - 1u);

    SYS_HANDLE currentDirectoryHandle = SYS_HANDLE_INVALID;
    KSTATUS currentDirectoryStatus = NtlOpenDir(currentDirectory, SYS_DIR_ACCESS_OPEN,
                                                SYS_DIR_FLAG_NONE, &currentDirectoryHandle);

    SYS_HANDLE imageDirectory = SYS_HANDLE_INVALID;
    int imageDirectoryBorrowed = 0;
    status = NtlOpenImageDirectory(request, &imageDirectory, &imageDirectoryBorrowed);
    if (!KSUCCESS(status)) {
        if (currentDirectoryHandle != SYS_HANDLE_INVALID)
            (void)SysCloseHandle(currentDirectoryHandle);
        goto fail;
    }

    SYS_HANDLE primaryModuleRoot = imageDirectory;
    SYS_STRING_VIEW primaryModuleDirectory = SysSvFromParts((const char*)0, 0);
    SYS_HANDLE secondaryModuleRoot = request->ModuleRootHandle;
    SYS_STRING_VIEW secondaryModuleDirectory = request->ModuleDirectory;
    SYS_HANDLE fallbackModuleRoot = SYS_HANDLE_INVALID;
    SYS_STRING_VIEW fallbackModuleDirectory = SysSvFromParts((const char*)0, 0);

    // Current-directory lookup is a search-path optimization, not a bootstrap dependency.
    if (KSUCCESS(currentDirectoryStatus)) {
        primaryModuleRoot = currentDirectoryHandle;
        secondaryModuleRoot = imageDirectory;
        secondaryModuleDirectory = SysSvFromParts((const char*)0, 0);
        fallbackModuleRoot = request->ModuleRootHandle;
        fallbackModuleDirectory = request->ModuleDirectory;
    }

    status = NtlpCreateModuleContext(state->ProcessHandle, primaryModuleRoot,
                                     primaryModuleDirectory, secondaryModuleRoot,
                                     secondaryModuleDirectory, fallbackModuleRoot,
                                     fallbackModuleDirectory, request->ResolveModuleName,
                                     request->ResolveContext, 0, &state->Modules);
    if (!KSUCCESS(status)) {
        if (!imageDirectoryBorrowed) (void)SysCloseHandle(imageDirectory);
        if (currentDirectoryHandle != SYS_HANDLE_INVALID)
            (void)SysCloseHandle(currentDirectoryHandle);
        goto fail;
    }

    memset(&state->Image, 0, sizeof(state->Image));
    status = NtlpLoadMainImage(state->Modules, imageFile, &state->Image);
    if (!imageDirectoryBorrowed) (void)SysCloseHandle(imageDirectory);
    if (currentDirectoryHandle != SYS_HANDLE_INVALID)
        (void)SysCloseHandle(currentDirectoryHandle);
    if (!KSUCCESS(status)) goto fail;

    memset(state->ChildCapabilities, 0, sizeof(state->ChildCapabilities));
    for (uint32_t i = 0; i < request->StartupCapabilityCount; ++i) {
        const NTL_PROCESS_STARTUP_CAPABILITY* source = &request->StartupCapabilities[i];
        if (!SysStartupCapabilityNameValid(source->Name) || source->SourceHandle == SYS_HANDLE_INVALID) {
            status = NTLSTATUS_INVALID_PARAMETER;
            goto fail;
        }
        for (uint32_t j = 0; j < i; ++j) {
            if (strcmp(state->ChildCapabilities[j].Name, source->Name) == 0) {
                status = NTLSTATUS_INVALID_PARAMETER;
                goto fail;
            }
        }
        SYS_HANDLE childHandle = SYS_HANDLE_INVALID;
        SYS_DUPLICATE_HANDLE_REQUEST duplicate;
        memset(&duplicate, 0, sizeof(duplicate));
        duplicate.Size = sizeof(duplicate);
        duplicate.SourceHandle = source->SourceHandle;
        duplicate.TargetProcessHandle = state->ProcessHandle;
        duplicate.DesiredRights = source->DesiredRights;
        duplicate.Flags = source->DuplicateFlags;
        duplicate.TargetHandleOut = &childHandle;
        status = SysDuplicateHandle(&duplicate);
        if (!KSUCCESS(status)) goto fail;
        memcpy(state->ChildCapabilities[i].Name, source->Name, sizeof(source->Name));
        state->ChildCapabilities[i].Handle = childHandle;
    }
    state->StartupCapabilityCount = request->StartupCapabilityCount;

    status = NtlpBuildLoaderData(state->Modules, currentDirectory, request->Arguments,
                                 request->Environment, request->StartupData,
                                 request->StartupDataSize, &state->LoaderData);
    if (!KSUCCESS(status)) goto fail;

    *preparationOut = state;
    status = NTLSTATUS_SUCCESS;
    goto done;

fail:
    NtlDestroyPreparation(state, true);
done:
    if (openedImage != SYS_HANDLE_INVALID) (void)SysCloseHandle(openedImage);
    return status;
}

// Returns the borrowed process capability of a prepared child.
KSTATUS NtlPreparedProcessGetHandle(const NTL_PROCESS_PREPARATION* preparation,
                                    SYS_HANDLE* processHandleOut) {
    if (!processHandleOut) return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = SYS_HANDLE_INVALID;
    if (!preparation || preparation->ProcessHandle == SYS_HANDLE_INVALID ||
        !preparation->Modules || !preparation->LoaderData)
        return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = preparation->ProcessHandle;
    return NTLSTATUS_SUCCESS;
}

// Starts and consumes one prepared child, including on ordinary start failure.
KSTATUS NtlStartPreparedProcess(NTL_PROCESS_PREPARATION* preparation,
                                SYS_HANDLE* processHandleOut) {
    if (!preparation || !processHandleOut) return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = SYS_HANDLE_INVALID;
    if (preparation->ProcessHandle == SYS_HANDLE_INVALID ||
        !preparation->Modules || !preparation->LoaderData)
        return NTLSTATUS_INVALID_PARAMETER;

    SYS_START_PROCESS_REQUEST startRequest;
    memset(&startRequest, 0, sizeof(startRequest));
    startRequest.Size = sizeof(startRequest);
    startRequest.ProcessHandle = preparation->ProcessHandle;
    startRequest.Flags = SYS_START_PROCESS_FLAG_NONE;
    startRequest.MainImageBase = preparation->Image.Base;
    startRequest.MainImageSize = preparation->Image.Size;
    startRequest.EntryPoint = preparation->Image.EntryPoint;
    startRequest.StartupCapabilities = preparation->ChildCapabilities;
    startRequest.StartupCapabilityCount = preparation->StartupCapabilityCount;
    startRequest.LoaderData = preparation->LoaderData;
    KSTATUS status = SysStartProcess(&startRequest);
    if (!KSUCCESS(status)) {
        NtlDestroyPreparation(preparation, true);
        return status;
    }

    // Transfer the native handle before releasing loader-owned metadata.
    *processHandleOut = preparation->ProcessHandle;
    preparation->ProcessHandle = SYS_HANDLE_INVALID;
    NtlDestroyPreparation(preparation, false);
    return NTLSTATUS_SUCCESS;
}

static void NtlPreparationRequestFromCreate(const NTL_CREATE_PROCESS_AT_REQUEST* create,
                                             NTL_PREPARE_PROCESS_REQUEST* preparedRequest) {
    memset(preparedRequest, 0, sizeof(*preparedRequest));
    preparedRequest->Size = sizeof(*preparedRequest);
    preparedRequest->Flags = create->Flags;
    preparedRequest->RootHandle = create->RootHandle;
    preparedRequest->Path = create->Path;
    preparedRequest->CommandLine = create->CommandLine;
    preparedRequest->CurrentDirectory = create->CurrentDirectory;
    preparedRequest->SecurityContext = create->SecurityContext;
    preparedRequest->ModuleRootHandle = create->ModuleRootHandle;
    preparedRequest->ModuleDirectory = create->ModuleDirectory;
    preparedRequest->ResolveModuleName = create->ResolveModuleName;
    preparedRequest->ResolveContext = create->ResolveContext;
    preparedRequest->StartupCapabilities = create->StartupCapabilities;
    preparedRequest->StartupCapabilityCount = create->StartupCapabilityCount;
    preparedRequest->Arguments = create->Arguments;
    preparedRequest->Environment = create->Environment;
    preparedRequest->StartupData = create->StartupData;
    preparedRequest->StartupDataSize = create->StartupDataSize;
}

static KSTATUS NtlCreateProcessFile(const NTL_CREATE_PROCESS_AT_REQUEST* request, SYS_HANDLE fileHandle) {
    NTL_PREPARE_PROCESS_REQUEST preparationRequest;
    NtlPreparationRequestFromCreate(request, &preparationRequest);
    NTL_PROCESS_PREPARATION* preparation = 0;
    KSTATUS status = NtlPrepareProcess(&preparationRequest, fileHandle, &preparation);
    if (!KSUCCESS(status)) return status;
    return NtlStartPreparedProcess(preparation, request->ProcessHandleOut);
}

// Creates and starts one child from a PE file beneath a capability-relative root.
KSTATUS NtlCreateProcessAt(const NTL_CREATE_PROCESS_AT_REQUEST* request) {
    if (!request || request->Size != sizeof(*request) ||
        request->Flags != NTL_PROCESS_LOAD_FLAG_NONE ||
        request->RootHandle == SYS_HANDLE_INVALID ||
        request->ModuleRootHandle == SYS_HANDLE_INVALID ||
        !request->Path.Data || request->Path.Length == 0 ||
        !SysSvIsValid(request->ModuleDirectory) ||
        request->StartupCapabilityCount > SYS_STARTUP_CAPABILITY_MAX ||
        (request->StartupCapabilityCount && !request->StartupCapabilities) ||
        (request->Arguments.Count && !request->Arguments.Entries) ||
        (request->Environment.Count && !request->Environment.Entries) ||
        request->StartupDataSize > NTL_PROCESS_STARTUP_DATA_MAX ||
        (request->StartupDataSize && !request->StartupData) ||
        !request->ProcessHandleOut)
        return NTLSTATUS_INVALID_PARAMETER;

    *request->ProcessHandleOut = SYS_HANDLE_INVALID;

    SYS_HANDLE fileHandle = SYS_HANDLE_INVALID;
    SYS_OPEN_FILE_AT_REQUEST open;
    memset(&open, 0, sizeof(open));
    open.Size = sizeof(open);
    open.RootHandle = request->RootHandle;
    open.Path = request->Path;
    open.Access = SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY;
    open.Disposition = SYS_FILE_OPEN_EXISTING;
    open.Flags = SYS_FILE_FLAG_NONE;
    open.HandleOut = &fileHandle;
    KSTATUS status = SysOpenFileAt(&open);
    if (!KSUCCESS(status))
        return status;

    status = NtlCreateProcessFile(request, fileHandle);
    (void)SysCloseHandle(fileHandle);
    return status;
}

KSTATUS NtlCreateSystemProcess(const NTL_CREATE_SYSTEM_PROCESS_REQUEST* request) {
    if (!request || request->Size != sizeof(*request) ||
        request->Flags != NTL_PROCESS_LOAD_FLAG_NONE || !request->Path.Data ||
        request->Path.Length == 0 || request->StartupDataSize > NTL_PROCESS_STARTUP_DATA_MAX ||
        (request->Arguments.Count && !request->Arguments.Entries) ||
        (request->Environment.Count && !request->Environment.Entries) ||
        (request->StartupDataSize && !request->StartupData) || !request->ProcessHandleOut)
        return NTLSTATUS_INVALID_PARAMETER;

    *request->ProcessHandleOut = SYS_HANDLE_INVALID;

    char canonicalPath[SYS_PATH_MAX + 1u];
    size_t canonicalSize = 0;
    KSTATUS canonicalStatus = NtlRealPath(request->Path, canonicalPath, sizeof(canonicalPath),
                                          &canonicalSize);
    SYS_STRING_VIEW imagePath = request->Path;
    if (KSUCCESS(canonicalStatus))
        imagePath = SysSvFromParts(canonicalPath, canonicalSize - 1u);
    else if (request->ImageFile == SYS_HANDLE_INVALID)
        return canonicalStatus;

    SYS_HANDLE systemRoot = SYS_HANDLE_INVALID;
    KSTATUS status = NtlResourceOpenSystemRoot(SYS_RIGHT_QUERY | SYS_VOLUME_RIGHT_OPEN, &systemRoot);
    if (!KSUCCESS(status)) return status;

    NTL_CREATE_PROCESS_AT_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Flags = request->Flags;
    create.RootHandle = systemRoot;
    create.Path = imagePath;
    create.CommandLine = request->CommandLine;
    create.CurrentDirectory = request->CurrentDirectory;
    create.SecurityContext = request->SecurityContext;
    create.ModuleRootHandle = systemRoot;
    create.ModuleDirectory = SysSvFromCstr("system");
    create.ResolveModuleName = request->ResolveModuleName;
    create.ResolveContext = request->ResolveContext;
    create.StartupCapabilities = request->StartupCapabilities;
    create.StartupCapabilityCount = request->StartupCapabilityCount;
    create.Arguments = request->Arguments;
    create.Environment = request->Environment;
    create.StartupData = request->StartupData;
    create.StartupDataSize = request->StartupDataSize;
    create.ProcessHandleOut = request->ProcessHandleOut;
    if (request->ImageFile != SYS_HANDLE_INVALID) {
        status = NtlCreateProcessFile(&create, request->ImageFile);
    } else {
        SYS_HANDLE imageFile = SYS_HANDLE_INVALID;
        status = NtlOpenFile(create.Path, SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY,
                             SYS_FILE_OPEN_EXISTING, SYS_FILE_FLAG_NONE, &imageFile);
        if (KSUCCESS(status)) {
            status = NtlCreateProcessFile(&create, imageFile);
            (void)SysCloseHandle(imageFile);
        }
    }
    (void)SysCloseHandle(systemRoot);
    return status;
}