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

#pragma once

#include "../zen.h"

#include <filesystem>

namespace zen {

class BenchPurgeSubCmd : public ZenSubCmdBase
{
public:
	BenchPurgeSubCmd();
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	bool m_SingleProcess = false;
};

class BenchHttpSubCmd : public ZenSubCmdBase
{
public:
	BenchHttpSubCmd();
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	void RunFixedCount(const std::string& BaseUri, const std::string& Path);
	void RunContinuous(const std::string& BaseUri, const std::string& Path);

	std::string m_Url;
	std::string m_SocketPath;
	int			m_Count		  = 100;
	int			m_Concurrency = 1;
	std::string m_Method	  = "GET";
	bool		m_NoKeepAlive = false;
	bool		m_Continuous  = false;
};

class BenchDiskSubCmd : public ZenSubCmdBase
{
public:
	BenchDiskSubCmd();
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	struct SeqResult
	{
		uint64_t BytesPerSec  = 0;
		uint64_t Iops		  = 0;
		double	 ElapsedSec	  = 0.0;
		int		 ErrorCount	  = 0;
		bool	 UsedDirectIo = false;
	};

	SeqResult RunSeqWrite(const std::filesystem::path& Dir, uint64_t BlockSize, bool DirectIo);
	SeqResult RunSeqRead(const std::filesystem::path& Dir, uint64_t BlockSize, bool DirectIo);
	SeqResult RunRandWrite(const std::filesystem::path& Dir, uint64_t BlockSize, bool DirectIo);
	SeqResult RunRandRead(const std::filesystem::path& Dir, uint64_t BlockSize, bool DirectIo);
	void	  PrefillRandFiles(const std::filesystem::path& Dir, bool DirectIo);
	void	  RunFileOps(const std::filesystem::path& Dir);
	void	  RunClone(const std::filesystem::path& Dir);
	void	  CleanupSeqFiles(const std::filesystem::path& Dir);
	void	  CleanupRandFiles(const std::filesystem::path& Dir);
	void	  CleanupFileOpFiles(const std::filesystem::path& Dir);
	void	  CleanupCloneFiles(const std::filesystem::path& Dir);
	void	  RunSync(const std::filesystem::path& Dir);
	void	  CleanupSyncFiles(const std::filesystem::path& Dir);

	std::string m_Path;
	std::string m_Run			  = "seq,rand,ops,clone";
	int			m_SyncDurationSec = 10;
	int			m_Concurrency	  = 1;
	uint64_t	m_FileSize		  = 256 * 1024 * 1024;
	int			m_FileCount		  = 1000;
	int			m_RandOps		  = 10000;
	bool		m_NoDirectIo	  = false;
};

class BenchCommand : public ZenCmdWithSubCommands
{
public:
	static constexpr char Name[]		= "bench";
	static constexpr char Description[] = "Utility command for benchmarking";

	BenchCommand();
	~BenchCommand();

	cxxopts::Options& Options() override { return m_Options; }
	ZenCmdCategory&	  CommandCategory() const override { return g_UtilitiesCategory; }

private:
	cxxopts::Options m_Options{Name, Description};
	std::string		 m_SubCommand;
	BenchPurgeSubCmd m_PurgeSubCmd;
	BenchHttpSubCmd	 m_HttpSubCmd;
	BenchDiskSubCmd	 m_DiskSubCmd;
};

}  // namespace zen