// POSIX borrowed-child substrate used by the later public vfork continuation layer.

#include <stdlib.h>
#include <string.h>
#include <ntlite/status.h>
#include <ntlite/syscall.h>
#include "posix_internal.h"
#include "posix_fd.h"

// Allocates one borrowed logical child and pins its inherited descriptor snapshot.
KSTATUS PosixVforkBeginBorrowedChild(POSIX_VFORK_CHILD_STATE** childOut) {
    if (!childOut) return NTLSTATUS_INVALID_PARAMETER;
    *childOut = (POSIX_VFORK_CHILD_STATE*)0;

    POSIX_THREAD_STATE* parentThread = PosixRuntimeThreadState();
    POSIX_PROCESS_STATE* nativeProcess = PosixRuntimeProcessState();
    if (!parentThread || !nativeProcess ||
        parentThread->LogicalProcess != &nativeProcess->DefaultIdentity)
        return NTLSTATUS_NOT_SUPPORTED;

    POSIX_VFORK_CHILD_STATE* child =
        (POSIX_VFORK_CHILD_STATE*)calloc(1u, sizeof(*child));
    if (!child) return NTLSTATUS_NO_MEMORY;
    child->LogicalProcess.FdLock = SYS_HANDLE_INVALID;
    child->InitialSignalMask = parentThread->SignalBlocked &
                               NTL_POSIX_SIGNAL_VALID_MASK &
                               ~NTL_POSIX_SIGNAL_UNBLOCKABLE_MASK;
    memcpy(child->LogicalProcess.SignalActions,
           parentThread->LogicalProcess->SignalActions,
           sizeof(child->LogicalProcess.SignalActions));

    KSTATUS status = PosixFdCreateInheritanceSnapshot(0u, &child->FdInheritanceToken);
    if (!KSUCCESS(status)) {
        free(child);
        return status;
    }

    NTL_POSIX_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Version = NTL_POSIX_PROTOCOL_VERSION;
    request.Operation = NTL_POSIX_OP_VFORK_BEGIN;
    request.FdInheritanceToken = child->FdInheritanceToken;

    NTL_POSIX_REPLY reply;
    status = PosixRuntimeCallServer(&request, &reply);
    if (!KSUCCESS(status) || !KSUCCESS(reply.Status)) {
        KSTATUS failure = KSUCCESS(status) ? reply.Status : status;
        PosixFdCancelInheritanceSnapshot(child->FdInheritanceToken);
        free(child);
        return failure;
    }
    if (reply.SpawnedPosixPid == 0 || reply.VforkToken == 0 ||
        reply.SpawnedPosixPid > 0x7FFFFFFFu) {
        if (reply.SpawnedPosixPid != 0 && reply.VforkToken != 0) {
            child->LogicalProcess.PosixPid = reply.SpawnedPosixPid;
            child->LogicalProcess.Generation = 1u;
            child->ServerToken = reply.VforkToken;
            if (KSUCCESS(PosixVforkCancelBorrowedChild(child)))
                return NTLSTATUS_FAILURE;
        }
        free(child);
        return NTLSTATUS_FAILURE;
    }

    child->LogicalProcess.PosixPid = reply.SpawnedPosixPid;
    child->LogicalProcess.Generation = 1u;
    child->ServerToken = reply.VforkToken;
    *childOut = child;
    return NTLSTATUS_SUCCESS;
}

// Attaches the calling native thread to one prepared borrowed logical child.
KSTATUS PosixVforkAttachBorrowedChild(POSIX_VFORK_CHILD_STATE* child) {
    if (!child || child->LogicalProcess.PosixPid == 0 || child->ServerToken == 0 ||
        child->FdInheritanceToken == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    SYS_TEB* teb = SysGetTeb();
    POSIX_PROCESS_STATE* nativeProcess = PosixRuntimeProcessState();
    if (!teb || teb->Magic != SYS_TEB_MAGIC || !nativeProcess || teb->SubsystemData)
        return NTLSTATUS_ACCESS_DENIED;

    POSIX_THREAD_STATE* thread = (POSIX_THREAD_STATE*)calloc(1u, sizeof(*thread));
    if (!thread) return NTLSTATUS_NO_MEMORY;
    thread->Magic = POSIX_THREAD_STATE_MAGIC;
    thread->NativeProcess = nativeProcess;
    thread->LogicalProcess = &child->LogicalProcess;
    thread->BorrowedChild = child;
    thread->NativeThreadId = SysGetCurrentTid();
    thread->SignalBlocked = child->InitialSignalMask;
    if (thread->NativeThreadId == 0) {
        free(thread);
        return NTLSTATUS_FAILURE;
    }

    KSTATUS status = SysSetUserApcDispatcher((void*)SysUserApcDispatcher);
    if (!KSUCCESS(status)) {
        free(thread);
        return status;
    }

    SYS_HANDLE selfThread = SYS_HANDLE_INVALID;
    status = SysOpenCurrentThread(SYS_RIGHT_QUERY | SYS_RIGHT_WAIT | SYS_RIGHT_DERIVE |
                                  SYS_RIGHT_TRANSFER | SYS_THREAD_RIGHT_NOTIFY, &selfThread);
    if (!KSUCCESS(status)) {
        free(thread);
        return status;
    }

    SYS_PORT_HANDLE_TRANSFER transfer;
    memset(&transfer, 0, sizeof(transfer));
    transfer.Handle = selfThread;
    transfer.DesiredRights = SYS_RIGHT_QUERY | SYS_RIGHT_WAIT | SYS_RIGHT_DERIVE |
                             SYS_THREAD_RIGHT_NOTIFY;
    transfer.Flags = SYS_PORT_TRANSFER_FLAG_NONE;

    NTL_POSIX_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Version = NTL_POSIX_PROTOCOL_VERSION;
    request.Operation = NTL_POSIX_OP_VFORK_ATTACH;
    request.TargetPosixPid = child->LogicalProcess.PosixPid;
    request.VforkToken = child->ServerToken;
    request.SignalMask = thread->SignalBlocked;
    request.SignalApcRoutine = (uint64_t)(uintptr_t)PosixSignalApc;

    NTL_POSIX_REPLY reply;
    status = PosixRuntimeCallServerWithHandles(&request, &transfer, 1u, &reply);
    (void)SysCloseHandle(selfThread);
    if (!KSUCCESS(status) || !KSUCCESS(reply.Status)) {
        KSTATUS failure = KSUCCESS(status) ? reply.Status : status;
        thread->Magic = 0;
        free(thread);
        return failure;
    }
    if (reply.PosixPid != child->LogicalProcess.PosixPid ||
        reply.NativeGeneration != child->LogicalProcess.Generation ||
        reply.ProcessState != NTL_POSIX_PROCESS_STATE_VFORK_BORROWED ||
        reply.VforkToken != child->ServerToken) {
        thread->Magic = 0;
        free(thread);
        return NTLSTATUS_FAILURE;
    }

    thread->ServerRegistered = 1u;
    teb->SubsystemData = thread;
    status = PosixFdInitializeLogicalProcess(&child->LogicalProcess,
                                             child->FdInheritanceToken);
    if (!KSUCCESS(status)) {
        teb->SubsystemData = (void*)0;
        thread->ServerRegistered = 0;
        thread->BorrowedChild = (POSIX_VFORK_CHILD_STATE*)0;
        thread->Magic = 0;
        free(thread);
        return status;
    }
    child->FdInheritanceToken = 0;
    return NTLSTATUS_SUCCESS;
}

// Tears down the calling borrowed thread's local wrappers before its native thread exits.
void PosixVforkDetachBorrowedChildLocal(POSIX_VFORK_CHILD_STATE* child, bool releaseServerHolders) {
    if (!child) return;
    SYS_TEB* teb = SysGetTeb();
    if (!teb || teb->Magic != SYS_TEB_MAGIC || !teb->SubsystemData) return;
    POSIX_THREAD_STATE* thread = (POSIX_THREAD_STATE*)teb->SubsystemData;
    if (thread->Magic != POSIX_THREAD_STATE_MAGIC ||
        thread->LogicalProcess != &child->LogicalProcess)
        return;

    if (releaseServerHolders) PosixFdDestroyLogicalProcess(&child->LogicalProcess);
    else PosixFdDestroyLogicalProcessLocal(&child->LogicalProcess);
    teb->SubsystemData = (void*)0;
    thread->ServerRegistered = 0;
    thread->BorrowedChild = (POSIX_VFORK_CHILD_STATE*)0;
    thread->Magic = 0;
    free(thread);
}

// Cancels one borrowed child after setup failure or after its native thread has exited.
KSTATUS PosixVforkCancelBorrowedChild(POSIX_VFORK_CHILD_STATE* child) {
    if (!child || child->LogicalProcess.PosixPid == 0 || child->ServerToken == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    NTL_POSIX_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Version = NTL_POSIX_PROTOCOL_VERSION;
    request.Operation = NTL_POSIX_OP_VFORK_CANCEL;
    request.TargetPosixPid = child->LogicalProcess.PosixPid;
    request.VforkToken = child->ServerToken;

    NTL_POSIX_REPLY reply;
    KSTATUS status = PosixRuntimeCallServer(&request, &reply);
    if (!KSUCCESS(status)) return status;
    if (!KSUCCESS(reply.Status)) return reply.Status;
    free(child);
    return NTLSTATUS_SUCCESS;
}

// Returns the borrowed-child state attached to the calling thread, if any.
POSIX_VFORK_CHILD_STATE* PosixVforkCurrentBorrowedChild(void) {
    POSIX_THREAD_STATE* thread = PosixRuntimeThreadState();
    return thread ? thread->BorrowedChild : (POSIX_VFORK_CHILD_STATE*)0;
}

// Records a genuine POSIX exit before the borrowed native thread self-terminates.
KSTATUS PosixVforkRecordExit(POSIX_VFORK_CHILD_STATE* child, int statusCode) {
    if (!child || child->LogicalProcess.PosixPid == 0 || child->ServerToken == 0)
        return NTLSTATUS_INVALID_PARAMETER;
    NTL_POSIX_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Version = NTL_POSIX_PROTOCOL_VERSION;
    request.Operation = NTL_POSIX_OP_VFORK_EXIT;
    request.TargetPosixPid = child->LogicalProcess.PosixPid;
    request.VforkToken = child->ServerToken;
    request.VforkExitCode = statusCode;
    NTL_POSIX_REPLY reply;
    KSTATUS status = PosixRuntimeCallServer(&request, &reply);
    if (!KSUCCESS(status)) return status;
    return reply.Status;
}

#include "posix_vfork_asm.h"

_Static_assert(offsetof(POSIX_VFORK_CONTEXT, ReturnAddress) == 0, "vfork return offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, ParentRsp) == 8, "vfork rsp offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, Xmm) == 80, "vfork xmm offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, Mxcsr) == 240, "vfork mxcsr offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, Child) == 248, "vfork child offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, ThreadHandle) == 256, "vfork handle offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, PreviousSuspendCount) == 264, "vfork suspend offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, ChildAttached) == 268, "vfork attached offset");
_Static_assert(offsetof(POSIX_VFORK_CONTEXT, ChildAttachStatus) == 272, "vfork status offset");
_Static_assert(sizeof(POSIX_VFORK_CONTEXT) == 280, "vfork context size");

// Allocates the persistent continuation and creates its borrowed native thread suspended.
POSIX_VFORK_CONTEXT* PosixVforkPrepare(const POSIX_VFORK_CONTEXT* captured) {
    if (!captured) {
        errno = EINVAL;
        return (POSIX_VFORK_CONTEXT*)0;
    }
    POSIX_VFORK_CONTEXT* context = (POSIX_VFORK_CONTEXT*)calloc(1u, sizeof(*context));
    if (!context) {
        errno = ENOMEM;
        return (POSIX_VFORK_CONTEXT*)0;
    }
    memcpy(context, captured, offsetof(POSIX_VFORK_CONTEXT, Child));

    KSTATUS status = PosixVforkBeginBorrowedChild(&context->Child);
    if (!KSUCCESS(status)) {
        free(context);
        errno = PosixStatusToErrno(status);
        return (POSIX_VFORK_CONTEXT*)0;
    }
    context->ChildPid = context->Child->LogicalProcess.PosixPid;
    context->ThreadHandle = SYS_HANDLE_INVALID;
    context->ChildAttachStatus = NTLSTATUS_FAILURE;

    SYS_CREATE_THREAD_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Entry = (uintptr_t)PosixVforkChildBootstrap;
    create.Argument = (uintptr_t)context;
    create.Priority = 1u;
    create.Flags = SYS_CREATE_THREAD_FLAG_SUSPENDED;
    create.ThreadHandleOut = &context->ThreadHandle;
    status = SysCreateThread(&create);
    if (!KSUCCESS(status)) {
        (void)PosixVforkCancelBorrowedChild(context->Child);
        context->Child = (POSIX_VFORK_CHILD_STATE*)0;
        free(context);
        errno = PosixStatusToErrno(status);
        return (POSIX_VFORK_CONTEXT*)0;
    }
    return context;
}

// Attaches on the private initial stack before entering the borrowed parent continuation.
void PosixVforkChildBootstrap(uintptr_t argument) {
    POSIX_VFORK_CONTEXT* context = (POSIX_VFORK_CONTEXT*)argument;
    if (!context || !context->Child) {
        SysExitThread((int32_t)NTLSTATUS_INVALID_PARAMETER);
        return;
    }
    KSTATUS status = PosixVforkAttachBorrowedChild(context->Child);
    context->ChildAttachStatus = status;
    if (!KSUCCESS(status)) {
        SysExitThread((int32_t)status);
        return;
    }
    _ReadWriteBarrier();
    context->ChildAttached = 1u;
    _ReadWriteBarrier();
    PosixVforkEnterChild(context);
    SysExitThread((int32_t)NTLSTATUS_FAILURE);
}

// Closes parent-side construction state after the atomic resume/wait has made the stack safe.
int PosixVforkParentFinish(POSIX_VFORK_CONTEXT* context, KSTATUS waitStatus) {
    if (!context) {
        errno = EIO;
        return -1;
    }
    if (context->ThreadHandle != SYS_HANDLE_INVALID) {
        (void)SysCloseHandle(context->ThreadHandle);
        context->ThreadHandle = SYS_HANDLE_INVALID;
    }

    int result = -1;
    if (!KSUCCESS(waitStatus)) {
        errno = PosixStatusToErrno(waitStatus);
    } else if (context->PreviousSuspendCount != 1u) {
        errno = EIO;
    } else if (!context->ChildAttached) {
        KSTATUS attachStatus = context->ChildAttachStatus;
        if (context->Child) {
            KSTATUS cancelStatus = PosixVforkCancelBorrowedChild(context->Child);
            if (KSUCCESS(cancelStatus)) context->Child = (POSIX_VFORK_CHILD_STATE*)0;
        }
        errno = PosixStatusToErrno(KSUCCESS(attachStatus) ? NTLSTATUS_FAILURE : attachStatus);
    } else {
        result = (int)context->ChildPid;
    }

    if (context->Child) free(context->Child);
    free(context);
    return result;
}