// Copyright Epic Games, Inc. All Rights Reserved.

#include "exitwatcher.h"

#include <zencore/logging.h>

ZEN_THIRD_PARTY_INCLUDES_START

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#	include <asio/io_context.hpp>
#	include <asio/windows/object_handle.hpp>
#elif ZEN_PLATFORM_LINUX
#	include <sys/syscall.h>
#	include <sys/wait.h>
#	include <unistd.h>
#	include <asio/io_context.hpp>
#	include <asio/posix/stream_descriptor.hpp>

#	ifndef SYS_pidfd_open
#		define SYS_pidfd_open 434	// x86_64
#	endif
#elif ZEN_PLATFORM_MAC
#	include <sys/event.h>
#	include <sys/wait.h>
#	include <unistd.h>
#	include <asio/io_context.hpp>
#	include <asio/posix/stream_descriptor.hpp>
#endif

ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

// ============================================================================
// Linux: pidfd_open + stream_descriptor
// ============================================================================

#if ZEN_PLATFORM_LINUX

struct ProcessExitWatcher::Impl
{
	asio::io_context&								m_IoContext;
	std::unique_ptr<asio::posix::stream_descriptor> m_Descriptor;
	int												m_PidFd = -1;
	int												m_Pid	= 0;

	explicit Impl(asio::io_context& IoContext) : m_IoContext(IoContext) {}

	~Impl() { Cancel(); }

	void Watch(const ProcessHandle& Handle, std::function<void(int ExitCode)> OnExit)
	{
		m_Pid = Handle.Pid();

		// pidfd_open returns an fd that becomes readable when the process exits.
		// Available since Linux 5.3.
		m_PidFd = static_cast<int>(syscall(SYS_pidfd_open, m_Pid, 0));
		if (m_PidFd < 0)
		{
			ZEN_WARN("pidfd_open failed for pid {}: {}", m_Pid, strerror(errno));
			return;
		}

		m_Descriptor = std::make_unique<asio::posix::stream_descriptor>(m_IoContext, m_PidFd);

		m_Descriptor->async_wait(asio::posix::stream_descriptor::wait_read,
								 [this, Callback = std::move(OnExit)](const asio::error_code& Ec) {
									 if (Ec)
									 {
										 return;  // Cancelled or error
									 }

									 int ExitCode = -1;
									 int Status	  = 0;
									 // The pidfd told us the process exited. Reap it with waitpid.
									 if (waitpid(m_Pid, &Status, WNOHANG) > 0)
									 {
										 if (WIFEXITED(Status))
										 {
											 ExitCode = WEXITSTATUS(Status);
										 }
										 else if (WIFSIGNALED(Status))
										 {
											 constexpr int kSignalExitBase = 128;
											 ExitCode					   = kSignalExitBase + WTERMSIG(Status);
										 }
									 }

									 Callback(ExitCode);
								 });
	}

	void Cancel()
	{
		if (m_Descriptor)
		{
			asio::error_code Ec;
			m_Descriptor->cancel(Ec);
			m_Descriptor.reset();
			// stream_descriptor closes the fd on destruction, so don't close m_PidFd separately
			m_PidFd = -1;
		}
		else if (m_PidFd >= 0)
		{
			close(m_PidFd);
			m_PidFd = -1;
		}
	}
};

// ============================================================================
// Windows: object_handle::async_wait
// ============================================================================

#elif ZEN_PLATFORM_WINDOWS

struct ProcessExitWatcher::Impl
{
	asio::io_context&							  m_IoContext;
	std::unique_ptr<asio::windows::object_handle> m_ObjectHandle;
	void*										  m_DuplicatedHandle = nullptr;

	explicit Impl(asio::io_context& IoContext) : m_IoContext(IoContext) {}

	~Impl() { Cancel(); }

	void Watch(const ProcessHandle& Handle, std::function<void(int ExitCode)> OnExit)
	{
		// Duplicate the process handle so ASIO can take ownership independently
		HANDLE SourceHandle	  = static_cast<HANDLE>(Handle.Handle());
		HANDLE CurrentProcess = GetCurrentProcess();
		BOOL   Success		  = DuplicateHandle(CurrentProcess,
										SourceHandle,
										CurrentProcess,
										reinterpret_cast<LPHANDLE>(&m_DuplicatedHandle),
										SYNCHRONIZE | PROCESS_QUERY_INFORMATION,
										FALSE,
										0);

		if (!Success)
		{
			ZEN_WARN("DuplicateHandle failed for pid {}: {}", Handle.Pid(), GetLastError());
			return;
		}

		// object_handle takes ownership of the handle
		m_ObjectHandle = std::make_unique<asio::windows::object_handle>(m_IoContext, m_DuplicatedHandle);

		m_ObjectHandle->async_wait([this, DupHandle = m_DuplicatedHandle, Callback = std::move(OnExit)](const asio::error_code& Ec) {
			if (Ec)
			{
				return;
			}

			DWORD ExitCode = 0;
			GetExitCodeProcess(static_cast<HANDLE>(DupHandle), &ExitCode);
			Callback(static_cast<int>(ExitCode));
		});
	}

	void Cancel()
	{
		if (m_ObjectHandle)
		{
			asio::error_code Ec;
			m_ObjectHandle->cancel(Ec);
			m_ObjectHandle.reset();	 // Closes the duplicated handle
			m_DuplicatedHandle = nullptr;
		}
		else if (m_DuplicatedHandle)
		{
			CloseHandle(static_cast<HANDLE>(m_DuplicatedHandle));
			m_DuplicatedHandle = nullptr;
		}
	}
};

// ============================================================================
// macOS: kqueue EVFILT_PROC + stream_descriptor
// ============================================================================

#elif ZEN_PLATFORM_MAC

struct ProcessExitWatcher::Impl
{
	asio::io_context&								m_IoContext;
	std::unique_ptr<asio::posix::stream_descriptor> m_Descriptor;
	int												m_KqueueFd = -1;
	int												m_Pid	   = 0;

	explicit Impl(asio::io_context& IoContext) : m_IoContext(IoContext) {}

	~Impl() { Cancel(); }

	void Watch(const ProcessHandle& Handle, std::function<void(int ExitCode)> OnExit)
	{
		m_Pid = Handle.Pid();

		m_KqueueFd = kqueue();
		if (m_KqueueFd < 0)
		{
			ZEN_WARN("kqueue() failed for pid {}: {}", m_Pid, strerror(errno));
			return;
		}

		// Register interest in the process exit event
		struct kevent Change;
		EV_SET(&Change, static_cast<uintptr_t>(m_Pid), EVFILT_PROC, EV_ADD | EV_ONESHOT, NOTE_EXIT, 0, nullptr);

		if (kevent(m_KqueueFd, &Change, 1, nullptr, 0, nullptr) < 0)
		{
			ZEN_WARN("kevent register failed for pid {}: {}", m_Pid, strerror(errno));
			close(m_KqueueFd);
			m_KqueueFd = -1;
			return;
		}

		m_Descriptor = std::make_unique<asio::posix::stream_descriptor>(m_IoContext, m_KqueueFd);

		m_Descriptor->async_wait(asio::posix::stream_descriptor::wait_read,
								 [this, Callback = std::move(OnExit)](const asio::error_code& Ec) {
									 if (Ec)
									 {
										 return;
									 }

									 // Drain the kqueue event
									 struct kevent	 Event;
									 struct timespec Timeout = {0, 0};
									 kevent(m_KqueueFd, nullptr, 0, &Event, 1, &Timeout);

									 int ExitCode = -1;
									 int Status	  = 0;
									 if (waitpid(m_Pid, &Status, WNOHANG) > 0)
									 {
										 if (WIFEXITED(Status))
										 {
											 ExitCode = WEXITSTATUS(Status);
										 }
										 else if (WIFSIGNALED(Status))
										 {
											 constexpr int kSignalExitBase = 128;
											 ExitCode					   = kSignalExitBase + WTERMSIG(Status);
										 }
									 }

									 Callback(ExitCode);
								 });
	}

	void Cancel()
	{
		if (m_Descriptor)
		{
			asio::error_code Ec;
			m_Descriptor->cancel(Ec);
			m_Descriptor.reset();
			// stream_descriptor closes the kqueue fd on destruction
			m_KqueueFd = -1;
		}
		else if (m_KqueueFd >= 0)
		{
			close(m_KqueueFd);
			m_KqueueFd = -1;
		}
	}
};

#endif

// ============================================================================
// Common wrapper (delegates to Impl)
// ============================================================================

ProcessExitWatcher::ProcessExitWatcher(asio::io_context& IoContext) : m_Impl(std::make_unique<Impl>(IoContext))
{
}

ProcessExitWatcher::~ProcessExitWatcher() = default;

void
ProcessExitWatcher::Watch(const ProcessHandle& Handle, std::function<void(int ExitCode)> OnExit)
{
	m_Impl->Watch(Handle, std::move(OnExit));
}

void
ProcessExitWatcher::Cancel()
{
	m_Impl->Cancel();
}

}  // namespace zen