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

#pragma once

#include <zencompute/functionservice.h>
#include <zencompute/zencompute.h>
#include <zencore/basicfile.h>
#include <zencore/compactbinarybuilder.h>
#include <zenstore/cidstore.h>
#include <zenstore/gc.h>
#include <zenstore/zenstore.h>

#include <filesystem>
#include <functional>
#include <map>
#include <unordered_map>

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

#if ZEN_WITH_COMPUTE_SERVICES

namespace zen::compute {

//////////////////////////////////////////////////////////////////////////

struct RecordingFileWriter
{
	RecordingFileWriter(RecordingFileWriter&&) = delete;
	RecordingFileWriter& operator=(RecordingFileWriter&&) = delete;

	RwLock		   m_FileLock;
	BasicFile	   m_File;
	uint64_t	   m_FileOffset = 0;
	CbObjectWriter m_TocWriter;
	BasicFile	   m_TocFile;

	RecordingFileWriter();
	~RecordingFileWriter();

	void Open(std::filesystem::path FilePath);
	void Close();
	void AppendObject(const CbObject& Object, const IoHash& ObjectHash);
};

//////////////////////////////////////////////////////////////////////////

/**
 * Recording "runner" implementation
 *
 * This class writes out all actions and their attachments to a recording directory
 * in a format that can be read back by the RecordingReader.
 *
 * The contents of the recording directory will be self-contained, with all referenced
 * attachments stored in the recording directory itself, so that the recording can be
 * moved or shared without needing to maintain references to the main CID store.
 *
 */

class ActionRecorder
{
public:
	ActionRecorder(ChunkResolver& InChunkResolver, const std::filesystem::path& RecordingLogPath);
	~ActionRecorder();

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

	void Shutdown();
	void RegisterWorker(const CbPackage& WorkerPackage);
	bool RecordAction(Ref<RunnerAction> Action);

private:
	ChunkResolver&		  m_ChunkResolver;
	std::filesystem::path m_RecordingLogDir;

	RecordingFileWriter	  m_WorkersFile;
	RecordingFileWriter	  m_ActionsFile;
	GcManager			  m_Gc;
	CidStore			  m_CidStore{m_Gc};
	std::atomic<int>	  m_ChunkCounter{0};
	std::atomic<uint64_t> m_ChunkBytesCounter{0};
	std::atomic<int>	  m_ActionsCounter{0};
};

}  // namespace zen::compute

#endif	// ZEN_WITH_COMPUTE_SERVICES