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

#pragma once

#include "resourcemetrics.h"

#include <zenutil/zenserverprocess.h>

#include <atomic>
#include <filesystem>

namespace zen {

/**
 * Storage Server Instance
 *
 * This class manages the lifecycle of a storage server instance, and
 * provides functions to query its state. There should be one instance
 * per module ID.
 */
class StorageServerInstance
{
public:
	struct Configuration
	{
		uint16_t			  BasePort;
		std::filesystem::path HydrationTempPath;
		std::string			  HydrationTargetSpecification;
		uint32_t			  HttpThreadCount = 0;	// Automatic
		int					  CoreLimit		  = 0;	// Automatic
		std::filesystem::path ConfigPath;
	};

	StorageServerInstance(ZenServerEnvironment& RunEnvironment, const Configuration& Config, std::string_view ModuleId);
	~StorageServerInstance();

	const ResourceMetrics& GetResourceMetrics() const { return m_ResourceMetrics; }

	inline std::string_view GetModuleId() const { return m_ModuleId; }
	inline uint16_t			GetBasePort() const { return m_Config.BasePort; }
	void					GetProcessMetrics(ProcessMetrics& OutMetrics) const;

#if ZEN_PLATFORM_WINDOWS
	void SetJobObject(JobObject* InJobObject) { m_JobObject = InJobObject; }
#endif

	class SharedLockedPtr
	{
	public:
		SharedLockedPtr(const SharedLockedPtr&) = delete;
		SharedLockedPtr();

		SharedLockedPtr(RwLock& Lock, StorageServerInstance* Instance, bool Wait);

		SharedLockedPtr(SharedLockedPtr&& Rhs);
		~SharedLockedPtr();

		SharedLockedPtr& operator=(const SharedLockedPtr&) = delete;
		SharedLockedPtr& operator						   =(SharedLockedPtr&& Rhs);

		operator bool() const { return m_Instance != nullptr; }

		std::string_view GetModuleId() const;
		uint16_t		 GetBasePort() const
		{
			ZEN_ASSERT(m_Instance);
			return m_Instance->GetBasePort();
		}
		bool IsRunning() const;

		const ResourceMetrics& GetResourceMetrics() const
		{
			ZEN_ASSERT(m_Instance);
			return m_Instance->m_ResourceMetrics;
		}
		void UpdateMetrics()
		{
			ZEN_ASSERT(m_Instance);
			return m_Instance->UpdateMetricsLocked();
		}

#if ZEN_WITH_TESTS
		void TerminateForTesting() const;  // kills the child process to simulate a crash
#endif

	private:
		RwLock*				   m_Lock	  = nullptr;
		StorageServerInstance* m_Instance = nullptr;
	};

	[[nodiscard]] SharedLockedPtr LockShared(bool Wait) { return SharedLockedPtr(m_Lock, this, Wait); }

	class ExclusiveLockedPtr
	{
	public:
		ExclusiveLockedPtr(const ExclusiveLockedPtr&) = delete;
		ExclusiveLockedPtr();

		ExclusiveLockedPtr(RwLock& Lock, StorageServerInstance* Instance, bool Wait);

		ExclusiveLockedPtr(ExclusiveLockedPtr&& Rhs);
		~ExclusiveLockedPtr();

		ExclusiveLockedPtr& operator=(const ExclusiveLockedPtr&) = delete;
		ExclusiveLockedPtr& operator							 =(ExclusiveLockedPtr&& Rhs);

		operator bool() const { return m_Instance != nullptr; }

		std::string_view GetModuleId() const;
		uint16_t		 GetBasePort() const
		{
			ZEN_ASSERT(m_Instance);
			return m_Instance->GetBasePort();
		}
		bool IsRunning() const;

		const ResourceMetrics& GetResourceMetrics() const
		{
			ZEN_ASSERT(m_Instance);
			return m_Instance->m_ResourceMetrics;
		}

		void Provision();
		void Deprovision();
		void Hibernate();
		void Wake();

	private:
		RwLock*				   m_Lock	  = nullptr;
		StorageServerInstance* m_Instance = nullptr;
	};

	[[nodiscard]] ExclusiveLockedPtr LockExclusive(bool Wait) { return ExclusiveLockedPtr(m_Lock, this, Wait); }

private:
	void ProvisionLocked();
	void DeprovisionLocked();

	void HibernateLocked();
	void WakeLocked();

	void UpdateMetricsLocked();

	mutable RwLock		m_Lock;
	const Configuration m_Config;
	std::string			m_ModuleId;
	ZenServerInstance	m_ServerInstance;

	std::filesystem::path m_BaseDir;

	std::filesystem::path m_TempDir;
	ResourceMetrics		  m_ResourceMetrics;

	std::atomic<uint64_t> m_MemoryBytes		   = 0;
	std::atomic<uint64_t> m_KernelTimeMs	   = 0;
	std::atomic<uint64_t> m_UserTimeMs		   = 0;
	std::atomic<uint64_t> m_WorkingSetSize	   = 0;
	std::atomic<uint64_t> m_PeakWorkingSetSize = 0;
	std::atomic<uint64_t> m_PagefileUsage	   = 0;
	std::atomic<uint64_t> m_PeakPagefileUsage  = 0;

#if ZEN_PLATFORM_WINDOWS
	JobObject* m_JobObject = nullptr;
#endif

	void SpawnServerProcess();

	void Hydrate();
	void Dehydrate();

	friend class SharedLockedPtr;
	friend class ExclusiveLockedPtr;
};

}  // namespace zen