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

#pragma once

#include "../zen.h"

#include <zencompute/recordingreader.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/guid.h>
#include <zencore/iohash.h>

#include <filesystem>
#include <functional>
#include <memory>
#include <unordered_map>

namespace zen {
class CbPackage;
class CbObject;
struct IoHash;
class ChunkResolver;
}  // namespace zen

#if ZEN_WITH_COMPUTE_SERVICES

namespace zen::compute {
class ComputeServiceSession;
}

namespace zen {

class ExecCommand;

struct ExecFunctionDefinition
{
	std::string FunctionName;
	zen::Guid	FunctionVersion;
	zen::Guid	BuildSystemVersion;
	zen::IoHash WorkerId;
};

//////////////////////////////////////////////////////////////////////////
// Subcommands

class ExecHttpSubCmd : public ZenSubCmdBase
{
public:
	explicit ExecHttpSubCmd(ExecCommand& Parent);
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	ExecCommand& m_Parent;
	std::string	 m_HostName;
};

class ExecInprocSubCmd : public ZenSubCmdBase
{
public:
	explicit ExecInprocSubCmd(ExecCommand& Parent);
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	ExecCommand& m_Parent;
	bool		 m_Managed = false;
};

class ExecBeaconSubCmd : public ZenSubCmdBase
{
public:
	explicit ExecBeaconSubCmd(ExecCommand& Parent);
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	ExecCommand&		  m_Parent;
	std::string			  m_OrchestratorUrl;
	std::filesystem::path m_BeaconPath;
};

class ExecDumpSubCmd : public ZenSubCmdBase
{
public:
	explicit ExecDumpSubCmd(ExecCommand& Parent);
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	ExecCommand& m_Parent;
};

class ExecBuildlogSubCmd : public ZenSubCmdBase
{
public:
	explicit ExecBuildlogSubCmd(ExecCommand& Parent);
	void Run(const ZenCliOptions& GlobalOptions) override;

private:
	ExecCommand& m_Parent;
};

/**
 * Zen CLI command for executing functions from a recording
 *
 * Mostly for testing and debugging purposes
 */

class ExecCommand : public ZenCmdWithSubCommands
{
public:
	ExecCommand();
	~ExecCommand();

	static constexpr char Name[]		= "exec";
	static constexpr char Description[] = "Execute functions from a recording";

	cxxopts::Options& Options() override { return m_Options; }

	// Shared state & helpers (public for subcommand access)

	using FunctionDefinition = ExecFunctionDefinition;

	static void EmitFunctionList(const std::vector<FunctionDefinition>& FunctionList);
	void		RegisterWorkerFunctionsFromDescription(const zen::CbObject& WorkerDesc, const zen::IoHash& WorkerId);

	int RunSession(zen::compute::ComputeServiceSession& ComputeSession, std::string_view OrchestratorUrl = {});

	std::unordered_map<zen::IoHash, zen::CbPackage> m_WorkerMap;
	std::vector<FunctionDefinition>					m_FunctionList;
	zen::ChunkResolver*								m_ChunkResolver	  = nullptr;
	zen::compute::RecordingReaderBase*				m_RecordingReader = nullptr;
	bool											m_VerboseLogging  = false;
	bool											m_QuietLogging	  = false;
	bool											m_DumpActions	  = false;
	std::filesystem::path							m_OutputPath;
	bool											m_Binary = false;
	std::filesystem::path							m_RecordingLogPath;

private:
	bool OnParentOptionsParsed(const ZenCliOptions& GlobalOptions) override;

	cxxopts::Options	  m_Options{Name, Description};
	std::string			  m_SubCommand;
	std::filesystem::path m_RecordingPath;
	int					  m_Offset = 0;
	int					  m_Stride = 1;
	int					  m_Limit  = 0;
	bool				  m_Quiet  = false;

	std::unique_ptr<zen::compute::RecordingReader>	 m_Reader;
	std::unique_ptr<zen::compute::UeRecordingReader> m_UeReader;

	ExecHttpSubCmd	   m_HttpSubCmd{*this};
	ExecInprocSubCmd   m_InprocSubCmd{*this};
	ExecBeaconSubCmd   m_BeaconSubCmd{*this};
	ExecDumpSubCmd	   m_DumpSubCmd{*this};
	ExecBuildlogSubCmd m_BuildlogSubCmd{*this};
};

}  // namespace zen

#endif	// ZEN_WITH_COMPUTE_SERVICES