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

#pragma once

#include <zencore/zencore.h>

#include <zencore/thread.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>

#	include <memory>
#	include <system_error>

namespace zen {

/**
 * @brief Abstract base class for I/O thread pools used by the http.sys server.
 *
 * Two implementations are available:
 * - WinTpIoThreadPool: Uses the Windows Thread Pool API (default, current behavior)
 * - ExplicitIoThreadPool: Uses raw IOCP + std::thread with load-based scaling
 */
class WinIoThreadPool
{
public:
	virtual ~WinIoThreadPool() = default;

	/**
	 * @brief Bind an I/O handle to this pool with the given callback
	 */
	virtual void CreateIocp(HANDLE IoHandle, PTP_WIN32_IO_CALLBACK Callback, void* Context, std::error_code& ErrorCode) = 0;

	/**
	 * @brief Called before issuing an async I/O operation.
	 * The Windows TP implementation calls StartThreadpoolIo; the explicit implementation is a no-op.
	 */
	virtual void StartIo() = 0;

	/**
	 * @brief Called when an async I/O operation fails synchronously (to undo StartIo).
	 * The Windows TP implementation calls CancelThreadpoolIo; the explicit implementation is a no-op.
	 */
	virtual void CancelIo() = 0;

	/**
	 * @brief Factory method to create an I/O thread pool
	 * @param UseExplicitThreads If true, creates an ExplicitIoThreadPool; otherwise creates a WinTpIoThreadPool
	 * @param MinThreads Minimum number of threads
	 * @param MaxThreads Maximum number of threads
	 */
	static std::unique_ptr<WinIoThreadPool> Create(bool UseExplicitThreads, int MinThreads, int MaxThreads);
};

/**
 * @brief Windows Thread Pool implementation (wraps CreateThreadpool/CreateThreadpoolIo)
 */
class WinTpIoThreadPool : public WinIoThreadPool
{
public:
	WinTpIoThreadPool(int InThreadCount, int InMaxThreadCount);
	~WinTpIoThreadPool();

	virtual void CreateIocp(HANDLE IoHandle, PTP_WIN32_IO_CALLBACK Callback, void* Context, std::error_code& ErrorCode) override;
	virtual void StartIo() override;
	virtual void CancelIo() override;

private:
	PTP_POOL			m_ThreadPool   = nullptr;
	PTP_CLEANUP_GROUP	m_CleanupGroup = nullptr;
	PTP_IO				m_ThreadPoolIo = nullptr;
	TP_CALLBACK_ENVIRON m_CallbackEnvironment;
};

/**
 * @brief Explicit IOCP + std::thread implementation with load-based thread scaling
 *
 * Creates a raw I/O completion port and manages worker threads directly. Threads
 * scale up when all are busy servicing completions, and scale down when idle for
 * an extended period.
 */
class ExplicitIoThreadPool : public WinIoThreadPool
{
public:
	ExplicitIoThreadPool(int InMinThreadCount, int InMaxThreadCount);
	~ExplicitIoThreadPool();

	virtual void CreateIocp(HANDLE IoHandle, PTP_WIN32_IO_CALLBACK Callback, void* Context, std::error_code& ErrorCode) override;
	virtual void StartIo() override;
	virtual void CancelIo() override;

private:
	void WorkerThreadMain();
	void SpawnWorkerThread();

	HANDLE m_Iocp = nullptr;

	PTP_WIN32_IO_CALLBACK m_Callback = nullptr;
	void*				  m_Context	 = nullptr;

	int m_MinThreads;
	int m_MaxThreads;

	std::atomic<int>  m_TotalThreads{0};
	std::atomic<int>  m_ActiveCount{0};
	std::atomic<bool> m_ShuttingDown{false};

	RwLock					 m_ThreadListLock;
	std::vector<std::thread> m_Threads;
};

}  // namespace zen
#endif