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

#include "workthreadpool.h"

#include <zencore/logging.h>

namespace zen {

namespace detail {
	struct LambdaWork : IWork
	{
		LambdaWork(auto Work) : WorkFunction(Work) {}
		virtual void Execute() override { WorkFunction(); }

		std::function<void()> WorkFunction;
	};
}  // namespace detail

WorkerThreadPool::WorkerThreadPool(int InThreadCount)
{
	for (int i = 0; i < InThreadCount; ++i)
	{
		m_WorkerThreads.emplace_back(&WorkerThreadPool::WorkerThreadFunction, this);
	}
}

WorkerThreadPool::~WorkerThreadPool()
{
	m_WorkQueue.CompleteAdding();

	for (std::thread& Thread : m_WorkerThreads)
	{
		Thread.join();
	}

	m_WorkerThreads.clear();
}

void
WorkerThreadPool::ScheduleWork(Ref<IWork> Work)
{
	m_WorkQueue.Enqueue(std::move(Work));
}

void
WorkerThreadPool::ScheduleWork(std::function<void()>&& Work)
{
	m_WorkQueue.Enqueue(new detail::LambdaWork(Work));
}

void
WorkerThreadPool::WorkerThreadFunction()
{
	do
	{
		Ref<IWork> Work;
		if (m_WorkQueue.WaitAndDequeue(Work))
		{
			try
			{
				Work->Execute();
			}
			catch (std::exception& e)
			{
				Work->m_Exception = std::current_exception();

				ZEN_WARN("Caught exception in worker thread: {}", e.what());
			}
		}
		else
		{
			return;
		}
	} while (true);
}

}  // namespace zen