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

#include "zenutil/zenserverprocess.h"

#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/session.h>
#include <zencore/string.h>
#include <zencore/thread.h>

#include <atomic>

#if ZEN_PLATFORM_WINDOWS
#	include <atlbase.h>
#	include <zencore/windows.h>
#else
#	include <sys/mman.h>
#endif

//////////////////////////////////////////////////////////////////////////

namespace zen {

namespace zenutil {
#if ZEN_PLATFORM_WINDOWS
	class SecurityAttributes
	{
	public:
		inline SECURITY_ATTRIBUTES* Attributes() { return &m_Attributes; }

	protected:
		SECURITY_ATTRIBUTES m_Attributes{};
		SECURITY_DESCRIPTOR m_Sd{};
	};

	// Security attributes which allows any user access

	class AnyUserSecurityAttributes : public SecurityAttributes
	{
	public:
		AnyUserSecurityAttributes()
		{
			m_Attributes.nLength		= sizeof m_Attributes;
			m_Attributes.bInheritHandle = false;  // Disable inheritance

			const BOOL Success = InitializeSecurityDescriptor(&m_Sd, SECURITY_DESCRIPTOR_REVISION);

			if (Success)
			{
				if (!SetSecurityDescriptorDacl(&m_Sd, TRUE, (PACL)NULL, FALSE))
				{
					ThrowLastError("SetSecurityDescriptorDacl failed");
				}

				m_Attributes.lpSecurityDescriptor = &m_Sd;
			}
		}
	};
#endif	// ZEN_PLATFORM_WINDOWS

}  // namespace zenutil

//////////////////////////////////////////////////////////////////////////

ZenServerState::ZenServerState()
{
}

ZenServerState::~ZenServerState()
{
	if (m_OurEntry)
	{
		// Clean up our entry now that we're leaving

		m_OurEntry->Reset();
		m_OurEntry = nullptr;
	}

#if ZEN_PLATFORM_WINDOWS
	if (m_Data)
	{
		UnmapViewOfFile(m_Data);
	}

	if (m_hMapFile)
	{
		CloseHandle(m_hMapFile);
	}
#else
	if (m_Data != nullptr)
	{
		munmap(m_Data, m_MaxEntryCount * sizeof(ZenServerEntry));
	}

	int Fd = int(intptr_t(m_hMapFile));
	close(Fd);
#endif

	m_Data = nullptr;
}

void
ZenServerState::Initialize()
{
	size_t MapSize = m_MaxEntryCount * sizeof(ZenServerEntry);

#if ZEN_PLATFORM_WINDOWS
	// TODO: there's a small chance of a race here, this logic could be tightened up with a mutex to
	//		 ensure only a single process at a time creates the mapping
	// TODO: the fallback to Local instead of Global has a flaw where if you start a non-elevated instance
	//		 first then start an elevated instance second you'll have the first instance with a local
	//		 mapping and the second instance with a global mapping.  This kind of elevated/non-elevated
	//		 shouldn't be common, but handling for it should be improved in the future.

	HANDLE hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Global\\ZenMap");
	if (hMap == NULL)
	{
		hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Local\\ZenMap");
	}

	if (hMap == NULL)
	{
		// Security attributes to enable any user to access state
		zenutil::AnyUserSecurityAttributes Attrs;

		hMap = CreateFileMapping(INVALID_HANDLE_VALUE,	// use paging file
								 Attrs.Attributes(),	// allow anyone to access
								 PAGE_READWRITE,		// read/write access
								 0,						// maximum object size (high-order DWORD)
								 DWORD(MapSize),		// maximum object size (low-order DWORD)
								 L"Global\\ZenMap");	// name of mapping object

		if (hMap == NULL)
		{
			hMap = CreateFileMapping(INVALID_HANDLE_VALUE,						// use paging file
									 Attrs.Attributes(),						// allow anyone to access
									 PAGE_READWRITE,							// read/write access
									 0,											// maximum object size (high-order DWORD)
									 m_MaxEntryCount * sizeof(ZenServerEntry),	// maximum object size (low-order DWORD)
									 L"Local\\ZenMap");							// name of mapping object
		}

		if (hMap == NULL)
		{
			ThrowLastError("Could not open or create file mapping object for Zen server state");
		}
	}

	void* pBuf = MapViewOfFile(hMap,				 // handle to map object
							   FILE_MAP_ALL_ACCESS,	 // read/write permission
							   0,					 // offset high
							   0,					 // offset low
							   DWORD(MapSize));

	if (pBuf == NULL)
	{
		ThrowLastError("Could not map view of Zen server state");
	}
#else
	int Fd = shm_open("/UnrealEngineZen", O_RDWR | O_CREAT | O_CLOEXEC, 0666);
	if (Fd < 0)
	{
		ThrowLastError("Could not open a shared memory object");
	}
	fchmod(Fd, 0666);
	void* hMap = (void*)intptr_t(Fd);

	int Result = ftruncate(Fd, MapSize);
	ZEN_UNUSED(Result);

	void* pBuf = mmap(nullptr, MapSize, PROT_READ | PROT_WRITE, MAP_SHARED, Fd, 0);
	if (pBuf == MAP_FAILED)
	{
		ThrowLastError("Could not map view of Zen server state");
	}
#endif

	m_hMapFile	 = hMap;
	m_Data		 = reinterpret_cast<ZenServerEntry*>(pBuf);
	m_IsReadOnly = false;
}

bool
ZenServerState::InitializeReadOnly()
{
	size_t MapSize = m_MaxEntryCount * sizeof(ZenServerEntry);

#if ZEN_PLATFORM_WINDOWS
	HANDLE hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Global\\ZenMap");
	if (hMap == NULL)
	{
		hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Local\\ZenMap");
	}

	if (hMap == NULL)
	{
		return false;
	}

	void* pBuf = MapViewOfFile(hMap,		   // handle to map object
							   FILE_MAP_READ,  // read permission
							   0,			   // offset high
							   0,			   // offset low
							   MapSize);

	if (pBuf == NULL)
	{
		ThrowLastError("Could not map view of Zen server state");
	}
#else
	int Fd = shm_open("/UnrealEngineZen", O_RDONLY | O_CLOEXEC, 0666);
	if (Fd < 0)
	{
		return false;
	}
	void* hMap = (void*)intptr_t(Fd);

	void* pBuf = mmap(nullptr, MapSize, PROT_READ, MAP_PRIVATE, Fd, 0);
	if (pBuf == MAP_FAILED)
	{
		ThrowLastError("Could not map read-only view of Zen server state");
	}
#endif

	m_hMapFile = hMap;
	m_Data	   = reinterpret_cast<ZenServerEntry*>(pBuf);

	return true;
}

ZenServerState::ZenServerEntry*
ZenServerState::Lookup(int DesiredListenPort)
{
	for (int i = 0; i < m_MaxEntryCount; ++i)
	{
		if (m_Data[i].DesiredListenPort == DesiredListenPort)
		{
			return &m_Data[i];
		}
	}

	return nullptr;
}

ZenServerState::ZenServerEntry*
ZenServerState::Register(int DesiredListenPort)
{
	if (m_Data == nullptr)
	{
		return nullptr;
	}

	// Allocate an entry

	int Pid = GetCurrentProcessId();

	for (int i = 0; i < m_MaxEntryCount; ++i)
	{
		ZenServerEntry& Entry = m_Data[i];

		if (Entry.DesiredListenPort.load(std::memory_order_relaxed) == 0)
		{
			uint16_t Expected = 0;
			if (Entry.DesiredListenPort.compare_exchange_strong(Expected, uint16_t(DesiredListenPort)))
			{
				// Successfully allocated entry

				m_OurEntry = &Entry;

				Entry.Pid				  = Pid;
				Entry.EffectiveListenPort = 0;
				Entry.Flags				  = 0;

				const Oid SesId = GetSessionId();
				memcpy(Entry.SessionId, &SesId, sizeof SesId);

				return &Entry;
			}
		}
	}

	return nullptr;
}

void
ZenServerState::Sweep()
{
	if (m_Data == nullptr)
	{
		return;
	}

	ZEN_ASSERT(m_IsReadOnly == false);

	for (int i = 0; i < m_MaxEntryCount; ++i)
	{
		ZenServerEntry& Entry = m_Data[i];

		if (Entry.DesiredListenPort)
		{
			if (IsProcessRunning(Entry.Pid) == false)
			{
				ZEN_DEBUG("Sweep - pid {} not running, reclaiming entry (port {})", Entry.Pid, Entry.DesiredListenPort);

				Entry.Reset();
			}
		}
	}
}

void
ZenServerState::Snapshot(std::function<void(const ZenServerEntry&)>&& Callback)
{
	if (m_Data == nullptr)
	{
		return;
	}

	for (int i = 0; i < m_MaxEntryCount; ++i)
	{
		ZenServerEntry& Entry = m_Data[i];

		if (Entry.DesiredListenPort)
		{
			Callback(Entry);
		}
	}
}

void
ZenServerState::ZenServerEntry::Reset()
{
	Pid					= 0;
	DesiredListenPort	= 0;
	Flags				= 0;
	EffectiveListenPort = 0;
}

void
ZenServerState::ZenServerEntry::SignalShutdownRequest()
{
	Flags |= uint16_t(FlagsEnum::kShutdownPlease);
}

void
ZenServerState::ZenServerEntry::SignalReady()
{
	Flags |= uint16_t(FlagsEnum::kIsReady);
}

bool
ZenServerState::ZenServerEntry::AddSponsorProcess(uint32_t PidToAdd)
{
	for (std::atomic<uint32_t>& PidEntry : SponsorPids)
	{
		if (PidEntry.load(std::memory_order_relaxed) == 0)
		{
			uint32_t Expected = 0;
			if (PidEntry.compare_exchange_strong(Expected, PidToAdd))
			{
				// Success!
				return true;
			}
		}
		else if (PidEntry.load(std::memory_order_relaxed) == PidToAdd)
		{
			// Success, the because pid is already in the list
			return true;
		}
	}

	return false;
}

//////////////////////////////////////////////////////////////////////////

std::atomic<int> TestCounter{0};

ZenServerEnvironment::ZenServerEnvironment()
{
}

ZenServerEnvironment::~ZenServerEnvironment()
{
}

void
ZenServerEnvironment::Initialize(std::filesystem::path ProgramBaseDir)
{
	m_ProgramBaseDir = ProgramBaseDir;

	ZEN_DEBUG("Program base dir is '{}'", ProgramBaseDir);

	m_IsInitialized = true;
}

void
ZenServerEnvironment::InitializeForTest(std::filesystem::path ProgramBaseDir, std::filesystem::path TestBaseDir)
{
	m_ProgramBaseDir = ProgramBaseDir;
	m_TestBaseDir	 = TestBaseDir;

	ZEN_INFO("Program base dir is '{}'", ProgramBaseDir);
	ZEN_INFO("Cleaning test base dir '{}'", TestBaseDir);
	DeleteDirectories(TestBaseDir.c_str());

	m_IsTestInstance = true;
	m_IsInitialized	 = true;
}

std::filesystem::path
ZenServerEnvironment::CreateNewTestDir()
{
	using namespace std::literals;

	ExtendableWideStringBuilder<256> TestDir;
	TestDir << "test"sv << int64_t(++TestCounter);

	std::filesystem::path TestPath = m_TestBaseDir / TestDir.c_str();

	ZEN_INFO("Creating new test dir @ '{}'", TestPath);

	CreateDirectories(TestPath.c_str());

	return TestPath;
}

std::filesystem::path
ZenServerEnvironment::GetTestRootDir(std::string_view Path)
{
	std::filesystem::path Root = m_ProgramBaseDir.parent_path().parent_path();

	std::filesystem::path Relative{Path};

	return Root / Relative;
}

//////////////////////////////////////////////////////////////////////////

std::atomic<int> ChildIdCounter{0};

ZenServerInstance::ZenServerInstance(ZenServerEnvironment& TestEnvironment) : m_Env(TestEnvironment)
{
	ZEN_ASSERT(TestEnvironment.IsInitialized());
}

ZenServerInstance::~ZenServerInstance()
{
	Shutdown();
}

void
ZenServerInstance::SignalShutdown()
{
	m_ShutdownEvent.Set();
}

void
ZenServerInstance::Shutdown()
{
	if (m_Process.IsValid())
	{
		if (m_Terminate)
		{
			ZEN_INFO("Terminating zenserver process");
			m_Process.Terminate(111);
		}
		else
		{
			SignalShutdown();
			m_Process.Wait();
			m_Process.Reset();
		}
	}
}

void
ZenServerInstance::SpawnServer(int BasePort, std::string_view AdditionalServerArgs)
{
	ZEN_ASSERT(!m_Process.IsValid());  // Only spawn once

	const int MyPid	  = zen::GetCurrentProcessId();
	const int ChildId = ++ChildIdCounter;

	ExtendableStringBuilder<32> ChildEventName;
	ChildEventName << "Zen_Child_" << ChildId;
	NamedEvent ChildEvent{ChildEventName};

	CreateShutdownEvent(BasePort);

	ExtendableStringBuilder<32> LogId;
	LogId << "Zen" << ChildId;

	ExtendableStringBuilder<512> CommandLine;
	CommandLine << "zenserver" ZEN_EXE_SUFFIX_LITERAL;	// see CreateProc() call for actual binary path

	const bool IsTest = m_Env.IsTestEnvironment();

	if (IsTest)
	{
		if (!m_OwnerPid.has_value())
		{
			m_OwnerPid = MyPid;
		}

		CommandLine << " --test --log-id " << LogId;
	}

	if (m_OwnerPid.has_value())
	{
		CommandLine << " --owner-pid " << m_OwnerPid.value();
	}

	CommandLine << " --child-id " << ChildEventName;

	if (BasePort)
	{
		CommandLine << " --port " << BasePort;
		m_BasePort = BasePort;
	}

	if (!m_TestDir.empty())
	{
		CommandLine << " --data-dir ";
		PathToUtf8(m_TestDir.c_str(), CommandLine);
	}

	if (m_MeshEnabled)
	{
		CommandLine << " --mesh";
	}

	if (!AdditionalServerArgs.empty())
	{
		CommandLine << " " << AdditionalServerArgs;
	}

	std::filesystem::path CurrentDirectory = std::filesystem::current_path();

	ZEN_DEBUG("Spawning server '{}'", LogId);

	uint32_t CreationFlags = 0;
	if (!IsTest)
	{
		CreationFlags |= CreateProcOptions::Flag_NewConsole;
	}

	const std::filesystem::path BaseDir		  = m_Env.ProgramBaseDir();
	const std::filesystem::path Executable	  = BaseDir / "zenserver" ZEN_EXE_SUFFIX_LITERAL;
	CreateProcOptions			CreateOptions = {
		  .WorkingDirectory = &CurrentDirectory,
		  .Flags			= CreationFlags,
	  };
	CreateProcResult ChildPid = CreateProc(Executable, CommandLine.ToView(), CreateOptions);
#if ZEN_PLATFORM_WINDOWS
	if (!ChildPid && ::GetLastError() == ERROR_ELEVATION_REQUIRED)
	{
		ZEN_DEBUG("Regular spawn failed - spawning elevated server");
		CreateOptions.Flags |= CreateProcOptions::Flag_Elevated;
		ChildPid = CreateProc(Executable, CommandLine.ToView(), CreateOptions);
	}
#endif

	if (!ChildPid)
	{
		ThrowLastError("Server spawn failed");
	}

	ZEN_DEBUG("Server '{}' spawned OK", LogId);

	if (IsTest)
	{
		m_Process.Initialize(ChildPid);
	}

	m_ReadyEvent = std::move(ChildEvent);
}

void
ZenServerInstance::CreateShutdownEvent(int BasePort)
{
	ExtendableStringBuilder<32> ChildShutdownEventName;
	ChildShutdownEventName << "Zen_" << BasePort;
	ChildShutdownEventName << "_Shutdown";
	NamedEvent ChildShutdownEvent{ChildShutdownEventName};
	m_ShutdownEvent = std::move(ChildShutdownEvent);
}

void
ZenServerInstance::AttachToRunningServer(int BasePort)
{
	ZenServerState State;
	if (!State.InitializeReadOnly())
	{
		// TODO: return success/error code instead?
		throw std::runtime_error("No zen state found");
	}

	const ZenServerState::ZenServerEntry* Entry = nullptr;

	if (BasePort)
	{
		Entry = State.Lookup(BasePort);
	}
	else
	{
		State.Snapshot([&](const ZenServerState::ZenServerEntry& InEntry) { Entry = &InEntry; });
	}

	if (!Entry)
	{
		// TODO: return success/error code instead?
		throw std::runtime_error("No server found");
	}

	m_Process.Initialize(Entry->Pid);
	CreateShutdownEvent(BasePort);
}

void
ZenServerInstance::Detach()
{
	if (m_Process.IsValid())
	{
		m_Process.Reset();
		m_ShutdownEvent.Close();
	}
}

void
ZenServerInstance::WaitUntilReady()
{
	while (m_ReadyEvent.Wait(100) == false)
	{
		if (!m_Process.IsRunning() || !m_Process.IsValid())
		{
			ZEN_INFO("Wait abandoned by invalid process (running={})", m_Process.IsRunning());
			return;
		}
	}
}

bool
ZenServerInstance::WaitUntilReady(int Timeout)
{
	return m_ReadyEvent.Wait(Timeout);
}

std::string
ZenServerInstance::GetBaseUri() const
{
	ZEN_ASSERT(m_BasePort);

	return fmt::format("http://localhost:{}", m_BasePort);
}

}  // namespace zen