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

#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_PROCESS_BUILDER_MAGIC 0x3142504Eu

typedef struct NTL_PROCESS_BUILDER_INTERNAL {
    uint32_t Magic;
    uint32_t StartupCapabilityCount;
    SYS_HANDLE ProcessHandle;
    NTL_MODULE_CONTEXT* Modules;
    NTL_PE_IMAGE Image;
    void* LoaderData;
    SYS_STARTUP_CAPABILITY ChildCapabilities[SYS_STARTUP_CAPABILITY_MAX];
} NTL_PROCESS_BUILDER_INTERNAL;

typedef char NTL_PROCESS_BUILDER_SIZE_CHECK[
    sizeof(NTL_PROCESS_BUILDER_INTERNAL) <= sizeof(NTL_PROCESS_BUILDER) ? 1 : -1];

static NTL_PROCESS_BUILDER_INTERNAL* NtlBuilderInternal(NTL_PROCESS_BUILDER* builder) {
    return (NTL_PROCESS_BUILDER_INTERNAL*)(void*)builder;
}

static int NtlBuilderRequestValid(const NTL_PROCESS_BUILDER_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);
}

static void NtlBuilderReset(NTL_PROCESS_BUILDER* builder) {
    if (builder) memset(builder, 0, sizeof(*builder));
}

// Aborts one prepared inert child. This is deliberately safe on an empty builder.
void NtlProcessBuilderAbort(NTL_PROCESS_BUILDER* builder) {
    if (!builder) return;
    NTL_PROCESS_BUILDER_INTERNAL* state = NtlBuilderInternal(builder);
    if (state->Magic != NTL_PROCESS_BUILDER_MAGIC) {
        NtlBuilderReset(builder);
        return;
    }
    if (state->Modules) NtlpDestroyModuleContext(state->Modules);
    if (state->ProcessHandle != SYS_HANDLE_INVALID) NtlDiscardChild(state->ProcessHandle);
    NtlBuilderReset(builder);
}

// Opens the executable's containing directory as the private dependency-search root.
static KSTATUS NtlOpenImageDirectory(const NTL_PROCESS_BUILDER_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 NtlProcessBuilderPrepare(NTL_PROCESS_BUILDER* builder,
                                 const NTL_PROCESS_BUILDER_REQUEST* request,
                                 SYS_HANDLE imageFile) {
    if (!builder || !NtlBuilderRequestValid(request)) return NTLSTATUS_INVALID_PARAMETER;
    NtlBuilderReset(builder);
    NTL_PROCESS_BUILDER_INTERNAL* state = NtlBuilderInternal(builder);
    state->ProcessHandle = SYS_HANDLE_INVALID;

    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)) 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;
    KSTATUS status = SysCreateProcessObject(&create);
    if (!KSUCCESS(status)) goto done;

    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 (!source->Tag || source->SourceHandle == SYS_HANDLE_INVALID) {
            status = NTLSTATUS_INVALID_PARAMETER;
            goto fail;
        }
        for (uint32_t j = 0; j < i; ++j) {
            if (state->ChildCapabilities[j].Tag == source->Tag) {
                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;
        state->ChildCapabilities[i].Tag = source->Tag;
        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;

    state->Magic = NTL_PROCESS_BUILDER_MAGIC;
    status = NTLSTATUS_SUCCESS;
    goto done;

fail:
    if (state->Modules) NtlpDestroyModuleContext(state->Modules);
    if (state->ProcessHandle != SYS_HANDLE_INVALID) NtlDiscardChild(state->ProcessHandle);
    NtlBuilderReset(builder);
done:
    if (openedImage != SYS_HANDLE_INVALID) (void)SysCloseHandle(openedImage);
    return status;
}

// Returns the process capability owned by one prepared builder without transferring it.
KSTATUS NtlProcessBuilderGetProcessHandle(NTL_PROCESS_BUILDER* builder,
                                          SYS_HANDLE* processHandleOut) {
    if (!builder || !processHandleOut) return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = SYS_HANDLE_INVALID;
    NTL_PROCESS_BUILDER_INTERNAL* state = NtlBuilderInternal(builder);
    if (state->Magic != NTL_PROCESS_BUILDER_MAGIC ||
        state->ProcessHandle == SYS_HANDLE_INVALID || !state->Modules || !state->LoaderData)
        return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = state->ProcessHandle;
    return NTLSTATUS_SUCCESS;
}

// Starts one prepared child. On success ownership of the retained native process
// capability moves to the caller; on failure the child is destroyed.
KSTATUS NtlProcessBuilderStart(NTL_PROCESS_BUILDER* builder, SYS_HANDLE* processHandleOut) {
    if (!builder || !processHandleOut) return NTLSTATUS_INVALID_PARAMETER;
    *processHandleOut = SYS_HANDLE_INVALID;
    NTL_PROCESS_BUILDER_INTERNAL* state = NtlBuilderInternal(builder);
    if (state->Magic != NTL_PROCESS_BUILDER_MAGIC ||
        state->ProcessHandle == SYS_HANDLE_INVALID || !state->Modules || !state->LoaderData)
        return NTLSTATUS_INVALID_PARAMETER;

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

    *processHandleOut = state->ProcessHandle;
    NtlBuilderReset(builder);
    return NTLSTATUS_SUCCESS;
}

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

static KSTATUS NtlCreateProcessFile(const NTL_CREATE_PROCESS_AT_REQUEST* request, SYS_HANDLE fileHandle) {
    NTL_PROCESS_BUILDER_REQUEST builderRequest;
    NtlBuilderRequestFromCreate(request, &builderRequest);
    NTL_PROCESS_BUILDER builder;
    memset(&builder, 0, sizeof(builder));
    KSTATUS status = NtlProcessBuilderPrepare(&builder, &builderRequest, fileHandle);
    if (!KSUCCESS(status)) return status;
    status = NtlProcessBuilderStart(&builder, request->ProcessHandleOut);
    if (!KSUCCESS(status)) NtlProcessBuilderAbort(&builder);
    return status;
}

// 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;
}