// 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 {
	const int LargeWorkerThreadPoolTreadCount  = gsl::narrow<int>(Max(std::thread::hardware_concurrency() - 1u, 2u));
	const int MediumWorkerThreadPoolTreadCount = gsl::narrow<int>(Max((std::thread::hardware_concurrency() / 4u), 2u));
	const int SmallWorkerThreadPoolTreadCount  = gsl::narrow<int>(Max((std::thread::hardware_concurrency() / 8u), 1u));
	const int TinyWorkerThreadPoolTreadCount   = 1;

	static 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 = "LargeThreadPool(burst)"};
	WorkerPool BackgroundLargeWorkerPool = {.TreadCount = LargeWorkerThreadPoolTreadCount, .Name = "LargeThreadPool(bkg)"};

	WorkerPool BurstMediumWorkerPool	  = {.TreadCount = MediumWorkerThreadPoolTreadCount, .Name = "MediumThreadPool(burst)"};
	WorkerPool BackgroundMediumWorkerPool = {.TreadCount = MediumWorkerThreadPoolTreadCount, .Name = "MediumThreadPool(bkg)"};

	WorkerPool BurstSmallWorkerPool		 = {.TreadCount = SmallWorkerThreadPoolTreadCount, .Name = "SmallThreadPool(burst)"};
	WorkerPool BackgroundSmallWorkerPool = {.TreadCount = SmallWorkerThreadPoolTreadCount, .Name = "SmallThreadPool(bkg)"};

	WorkerPool BurstTinyWorkerPool		= {.TreadCount = TinyWorkerThreadPoolTreadCount, .Name = "TinyThreadPool(burst)"};
	WorkerPool BackgroundTinyWorkerPool = {.TreadCount = TinyWorkerThreadPoolTreadCount, .Name = "TinyThreadPool(bkg)"};

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

	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;
	}
}  // namespace

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

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

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

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

WorkerThreadPool&
GetSyncWorkerPool()
{
	return EnsurePoolPtr(SyncWorkerPool);
}

void
ShutdownWorkerPools()
{
	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();
}
}  // namespace zen