#include <rx/fork.hpp>

#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <mutex>
#include <windows.h>
#else
#include <unistd.h>
#endif

#ifdef _WIN32
namespace {
struct ClientId {
  HANDLE UniqueProcess;
  HANDLE UniqueThread;
};

struct SECTION_IMAGE_INFORMATION {
  PVOID TransferAddress; // The address of the image entry point function.
  ULONG ZeroBits; // The number of high-order address bits that must be zero in
                  // the image base address.
  SIZE_T MaximumStackSize; // The maximum stack size of threads from the PE file
                           // header.
  SIZE_T CommittedStackSize; // The initial stack size of threads from the PE
                             // file header.
  ULONG SubSystemType; // The image subsystem from the PE file header (e.g.,
                       // Windows GUI, Windows CUI, POSIX).
  union {
    struct {
      USHORT SubSystemMinorVersion;
      USHORT SubSystemMajorVersion;
    };
    ULONG SubSystemVersion;
  };
  union {
    struct {
      USHORT MajorOperatingSystemVersion;
      USHORT MinorOperatingSystemVersion;
    };
    ULONG OperatingSystemVersion;
  };
  USHORT ImageCharacteristics; // The image characteristics from the PE file
                               // header.
  USHORT DllCharacteristics;   // The DLL characteristics flags (e.g., ASLR, NX
                               // compatibility).
  USHORT Machine;              // The image architecture (e.g., x86, x64, ARM).
  BOOLEAN ImageContainsCode;   // The image contains native executable code.
  union {
    UCHAR ImageFlags;
    struct {
      UCHAR ComPlusNativeReady
          : 1; // The image contains precompiled .NET assembly generated by NGEN
               // (Native Image Generator).
      UCHAR ComPlusILOnly : 1; // the image contains only Microsoft Intermediate
                               // Language (IL) assembly.
      UCHAR ImageDynamicallyRelocated
          : 1; // The image was mapped using a random base address rather than
               // the preferred base address.
      UCHAR ImageMappedFlat
          : 1; // The image was mapped using a single contiguous region, rather
               // than separate regions for each section.
      UCHAR BaseBelow4gb : 1; // The image was mapped using a base address below
                              // the 4 GB boundary.
      UCHAR ComPlusPrefer32bit : 1; // The image prefers to run as a 32-bit
                                    // process, even on a 64-bit system.
      UCHAR Reserved : 2;
    };
  };
  ULONG LoaderFlags; // Reserved by ntdll.dll for the Windows loader.
  ULONG
  ImageFileSize;  // The size of the image, in bytes, including all headers.
  ULONG CheckSum; // The image file checksum, from the PE optional header.
};

struct RtlUserProcessInformation {
  ULONG Length;
  HANDLE Process;
  HANDLE Thread;
  ClientId ClientId;
  SECTION_IMAGE_INFORMATION ImageInformation;
};

#define RTL_CLONE_PROCESS_FLAGS_CREATE_SUSPENDED 0x00000001
#define RTL_CLONE_PROCESS_FLAGS_INHERIT_HANDLES 0x00000002
#define RTL_CLONE_PROCESS_FLAGS_NO_SYNCHRONIZE 0x00000004

using RtlCloneUserProcessFn = long(__stdcall *)(
    ULONG Flags, PSECURITY_DESCRIPTOR ProcessSecurityDescriptor,
    PSECURITY_DESCRIPTOR ThreadSecurityDescriptor, HANDLE DebugPort,
    RtlUserProcessInformation *ProcessInformation);

RtlCloneUserProcessFn getRtlClone() {
  static auto fn = []() -> RtlCloneUserProcessFn {
    if (auto mod = ::GetModuleHandleW(L"ntdll.dll")) {
      return reinterpret_cast<RtlCloneUserProcessFn>(
          ::GetProcAddress(mod, "RtlCloneUserProcess"));
    }

    return nullptr;
  }();

  return fn;
}
} // namespace
#endif

rx::ProcessId rx::fork() {
#ifdef _WIN32
  auto rtlClone = getRtlClone();
  if (!rtlClone) {
    return static_cast<ProcessId>(-1);
  }

  auto parentPid = ::GetCurrentProcessId();

  RtlUserProcessInformation procInfo{};
  procInfo.Length = sizeof(procInfo);

  ULONG rtlFlags = RTL_CLONE_PROCESS_FLAGS_CREATE_SUSPENDED |
                   RTL_CLONE_PROCESS_FLAGS_INHERIT_HANDLES;

  const long status = rtlClone(rtlFlags, nullptr, nullptr, nullptr, &procInfo);

  if (status < 0) {
    return static_cast<ProcessId>(-1);
  }

  if (status == 0) {
    // parent process
    const DWORD childPid = ::GetProcessId(procInfo.Process);

    if (procInfo.Thread) {
      ::ResumeThread(procInfo.Thread);
      ::CloseHandle(procInfo.Thread);
    }

    ::CloseHandle(procInfo.Process);
    return static_cast<ProcessId>(childPid);
  }

  FreeConsole();
  AttachConsole(parentPid);

  // child process
  return ProcessId{};
#else
  return static_cast<ProcessId>(::fork());
#endif
}