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

#pragma once

#include "hordeagentmessage.h"
#include "hordecomputesocket.h"

#include <zenhorde/hordeclient.h>

#include <zencore/logbase.h>

#include <filesystem>
#include <memory>
#include <string>

namespace zen::horde {

/** Manages the lifecycle of a single Horde compute agent.
 *
 *  Handles the full connection sequence for one provisioned machine:
 *  1. Connect via TCP transport (with optional AES encryption wrapping)
 *  2. Create a multiplexed ComputeSocket with agent (channel 0) and child (channel 100)
 *  3. Perform the Attach/Fork handshake to establish the child channel
 *  4. Upload zenserver binary via the WriteFiles/ReadBlob protocol
 *  5. Execute zenserver remotely via ExecuteV2
 *  6. Poll for ExecuteOutput (stdout) and ExecuteResult (exit code)
 */
class HordeAgent
{
public:
	explicit HordeAgent(const MachineInfo& Info);
	~HordeAgent();

	HordeAgent(const HordeAgent&) = delete;
	HordeAgent& operator=(const HordeAgent&) = delete;

	/** Perform the channel setup handshake (Attach on agent channel, Fork, Attach on child channel).
	 *  Returns false if the handshake times out or receives an unexpected message. */
	bool BeginCommunication();

	/** Upload binary files to the remote agent.
	 *  @param BundleDir     Directory containing .blob files.
	 *  @param BundleLocator Locator string identifying the bundle (from CreateBundle). */
	bool UploadBinaries(const std::filesystem::path& BundleDir, const std::string& BundleLocator);

	/** Execute a command on the remote machine. */
	void Execute(const char*		Exe,
				 const char* const* Args,
				 size_t				NumArgs,
				 const char*		WorkingDir = nullptr,
				 const char* const* EnvVars	   = nullptr,
				 size_t				NumEnvVars = 0,
				 bool				UseWine	   = false);

	/** Poll for output and results. Returns true if the agent is still running.
	 *  When LogOutput is true, remote stdout is logged via ZEN_INFO. */
	bool Poll(bool LogOutput = true);

	void CloseConnection();
	bool IsValid() const;

	const MachineInfo& GetMachineInfo() const { return m_MachineInfo; }

private:
	LoggerRef Log() { return m_Log; }

	std::unique_ptr<ComputeSocket>		 m_Socket;
	std::unique_ptr<AgentMessageChannel> m_AgentChannel;  ///< Channel 0: agent control
	std::unique_ptr<AgentMessageChannel> m_ChildChannel;  ///< Channel 100: child I/O

	LoggerRef	m_Log;
	bool		m_IsValid	= false;
	bool		m_HasErrors = false;
	MachineInfo m_MachineInfo;
};

}  // namespace zen::horde