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

#include "zenutil/workerpools.h"

#include <zencore/intmath.h>
#include <zencore/thread.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {
namespace {

	struct WorkerPools
	{
		const int LargeWorkerThreadPoolTreadCount  = gsl::narrow<int>(Max(GetHardwareConcurrency() - 1u, 2u));
		const int MediumWorkerThreadPoolTreadCount = gsl::narrow<int>(Max((GetHardwareConcurrency() / 4u), 2u));
		const int SmallWorkerThreadPoolTreadCount  = gsl::narrow<int>(Max((GetHardwareConcurrency() / 8u), 1u));
		const int TinyWorkerThreadPoolTreadCount   = 1;

		bool IsShutDown = false;

		RwLock PoolLock;

		struct WorkerPool
		{
			std::unique_ptr<WorkerThreadPool> Pool;
			const int						  TreadCount;
			const std::string_view			  Name;
		};

		WorkerPool BurstLargeWorkerPool		 = {.TreadCount = LargeWorkerThreadPoolTreadCount, .Name = "large"};
		WorkerPool BackgroundLargeWorkerPool = {.TreadCount = LargeWorkerThreadPoolTreadCount, .Name = "large_bg"};

		WorkerPool BurstMediumWorkerPool	  = {.TreadCount = MediumWorkerThreadPoolTreadCount, .Name = "medium"};
		WorkerPool BackgroundMediumWorkerPool = {.TreadCount = MediumWorkerThreadPoolTreadCount, .Name = "medium_bg"};

		WorkerPool BurstSmallWorkerPool		 = {.TreadCount = SmallWorkerThreadPoolTreadCount, .Name = "small"};
		WorkerPool BackgroundSmallWorkerPool = {.TreadCount = SmallWorkerThreadPoolTreadCount, .Name = "small_bg"};

		WorkerPool BurstTinyWorkerPool		= {.TreadCount = TinyWorkerThreadPoolTreadCount, .Name = "tiny"};
		WorkerPool BackgroundTinyWorkerPool = {.TreadCount = TinyWorkerThreadPoolTreadCount, .Name = "tiny_bg"};

		WorkerPool SyncWorkerPool = {.TreadCount = 0, .Name = "synctp"};

		WorkerThreadPool& EnsurePoolPtr(WorkerPool& Pool)
		{
			{
				RwLock::SharedLockScope _(PoolLock);
				if (Pool.Pool)
				{
					return *Pool.Pool;
				}
			}
			RwLock::ExclusiveLockScope _(PoolLock);
			ZEN_ASSERT(!IsShutDown);
			if (!Pool.Pool)
			{
				Pool.Pool.reset(new WorkerThreadPool(Pool.TreadCount, Pool.Name));
			}
			return *Pool.Pool;
		}

		void Shutdown()
		{
			RwLock::ExclusiveLockScope _(PoolLock);
			IsShutDown = true;
			BurstLargeWorkerPool.Pool.reset();
			BackgroundLargeWorkerPool.Pool.reset();
			BurstMediumWorkerPool.Pool.reset();
			BackgroundMediumWorkerPool.Pool.reset();
			BurstSmallWorkerPool.Pool.reset();
			BackgroundSmallWorkerPool.Pool.reset();
			BurstTinyWorkerPool.Pool.reset();
			BackgroundTinyWorkerPool.Pool.reset();
			SyncWorkerPool.Pool.reset();
		}
	};

	WorkerPools& Pools()
	{
		static WorkerPools _;
		return _;
	}

}  // namespace

WorkerThreadPool&
GetLargeWorkerPool(EWorkloadType WorkloadType)
{
	auto& p = Pools();
	return p.EnsurePoolPtr(WorkloadType == EWorkloadType::Burst ? p.BurstLargeWorkerPool : p.BackgroundLargeWorkerPool);
}

WorkerThreadPool&
GetMediumWorkerPool(EWorkloadType WorkloadType)
{
	auto& p = Pools();
	return p.EnsurePoolPtr(WorkloadType == EWorkloadType::Burst ? p.BurstMediumWorkerPool : p.BackgroundMediumWorkerPool);
}

WorkerThreadPool&
GetSmallWorkerPool(EWorkloadType WorkloadType)
{
	auto& p = Pools();
	return p.EnsurePoolPtr(WorkloadType == EWorkloadType::Burst ? p.BurstSmallWorkerPool : p.BackgroundSmallWorkerPool);
}

WorkerThreadPool&
GetTinyWorkerPool(EWorkloadType WorkloadType)
{
	auto& p = Pools();
	return p.EnsurePoolPtr(WorkloadType == EWorkloadType::Burst ? p.BurstTinyWorkerPool : p.BackgroundTinyWorkerPool);
}

WorkerThreadPool&
GetSyncWorkerPool()
{
	auto& p = Pools();
	return p.EnsurePoolPtr(p.SyncWorkerPool);
}

void
ShutdownWorkerPools()
{
	Pools().Shutdown();
}
}  // namespace zen