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

// Round-trip integration tests for BuildsOperationUploadFolder / BuildsOperationUpdateFolder.
// Runs in-process against CreateFileBuildStorage so no HTTP server is needed.

#include <zenremotestore/builds/buildupdatefolder.h>
#include <zenremotestore/builds/builduploadfolder.h>
#include <zenremotestore/builds/filebuildstorage.h>
#include <zenremotestore/chunking/chunkingcache.h>
#include <zenremotestore/chunking/chunkingcontroller.h>
#include <zenremotestore/transferthreadworkers.h>

#include <zencore/basicfile.h>
#include <zencore/compactbinary.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iohash.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/testing.h>
#include <zencore/testutils.h>
#include <zencore/workthreadpool.h>
#include <zenutil/progress.h>

#include <algorithm>
#include <atomic>
#include <filesystem>
#include <string>
#include <vector>

namespace zen {

void
buildoperations_tests_forcelink()
{
}

#if ZEN_WITH_TESTS

namespace buildops_test {

	using namespace std::literals;

	struct FolderSpec
	{
		uint64_t Seed				= 1;
		uint32_t SmallFileCount		= 40;
		uint32_t MediumFileCount	= 10;
		uint32_t LargeFileCount		= 2;
		uint32_t DuplicateFileCount = 6;
	};

	static IoBuffer MakeBlob(uint64_t Seed, size_t Size)
	{
		FastRandom Rnd{.Seed = Seed};
		IoBuffer   Blob(Size);
		uint8_t*   Data	  = static_cast<uint8_t*>(Blob.MutableData());
		size_t	   Offset = 0;
		while (Offset < Size)
		{
			uint64_t Word  = Rnd.Next();
			size_t	 Chunk = std::min<size_t>(sizeof(Word), Size - Offset);
			std::memcpy(Data + Offset, &Word, Chunk);
			Offset += Chunk;
		}
		return Blob;
	}

	static void WriteTestFile(const std::filesystem::path& Path, const IoBuffer& Blob)
	{
		CreateDirectories(Path.parent_path());
		zen::WriteFile(Path, Blob);
	}

	static std::vector<std::filesystem::path> MakeTestFolder(const std::filesystem::path& Root, const FolderSpec& Spec)
	{
		CreateDirectories(Root);
		FastRandom Rnd{.Seed = Spec.Seed};

		std::vector<std::filesystem::path> Written;
		Written.reserve(Spec.SmallFileCount + Spec.MediumFileCount + Spec.LargeFileCount + Spec.DuplicateFileCount);

		auto Emit = [&](std::string_view SubDir, uint32_t Index, size_t Size) {
			std::filesystem::path Rel = std::filesystem::path(std::string(SubDir)) / fmt::format("f_{:05}.bin", Index);
			WriteTestFile(Root / Rel, MakeBlob(Spec.Seed * 7919ull + Index, Size));
			Written.push_back(Rel);
		};

		for (uint32_t I = 0; I < Spec.SmallFileCount; ++I)
		{
			Emit("small", I, 1024u + static_cast<size_t>(Rnd.Next() & 0xFFFu));
		}
		for (uint32_t I = 0; I < Spec.MediumFileCount; ++I)
		{
			Emit("medium", I, 60u * 1024u + static_cast<size_t>(Rnd.Next() & 0x3FFFu));
		}
		for (uint32_t I = 0; I < Spec.LargeFileCount; ++I)
		{
			Emit("large", I, 900u * 1024u + static_cast<size_t>(Rnd.Next() & 0x1FFFFu));
		}

		// Duplicates of previously-written small files so upload can re-use blocks / chunks.
		for (uint32_t I = 0; I < Spec.DuplicateFileCount && Spec.SmallFileCount > 0; ++I)
		{
			std::filesystem::path Source = Root / Written[I % Spec.SmallFileCount];
			std::filesystem::path Rel	 = std::filesystem::path("dupes") / fmt::format("d_{:05}.bin", I);
			CreateDirectories((Root / Rel).parent_path());
			std::error_code Ec;
			std::filesystem::copy_file(Source, Root / Rel, std::filesystem::copy_options::overwrite_existing, Ec);
			if (!Ec)
			{
				Written.push_back(Rel);
			}
		}

		return Written;
	}

	static void CopyTreeExcludingZen(const std::filesystem::path& Src, const std::filesystem::path& Dst)
	{
		CreateDirectories(Dst);
		std::error_code Ec;
		for (auto It = std::filesystem::recursive_directory_iterator(Src, Ec); !Ec && It != std::filesystem::recursive_directory_iterator();
			 It.increment(Ec))
		{
			const std::filesystem::path Rel = std::filesystem::relative(It->path(), Src);
			if (!Rel.empty() && Rel.begin()->string() == ".zen")
			{
				It.disable_recursion_pending();
				continue;
			}
			if (It->is_directory())
			{
				CreateDirectories(Dst / Rel);
			}
			else if (It->is_regular_file())
			{
				CreateDirectories((Dst / Rel).parent_path());
				std::error_code CopyEc;
				std::filesystem::copy_file(It->path(), Dst / Rel, std::filesystem::copy_options::overwrite_existing, CopyEc);
			}
		}
	}

	static std::vector<std::filesystem::path> ListRelative(const std::filesystem::path& Root)
	{
		std::vector<std::filesystem::path> Paths;
		std::error_code					   Ec;
		for (auto It = std::filesystem::recursive_directory_iterator(Root, Ec);
			 !Ec && It != std::filesystem::recursive_directory_iterator();
			 It.increment(Ec))
		{
			const std::filesystem::path Rel = std::filesystem::relative(It->path(), Root);
			if (!Rel.empty() && Rel.begin()->string() == ".zen")
			{
				It.disable_recursion_pending();
				continue;
			}
			if (It->is_regular_file())
			{
				Paths.push_back(Rel);
			}
		}
		std::sort(Paths.begin(), Paths.end());
		return Paths;
	}

	static bool FoldersEquivalent(const std::filesystem::path& A, const std::filesystem::path& B)
	{
		const auto AFiles = ListRelative(A);
		const auto BFiles = ListRelative(B);
		if (AFiles != BFiles)
		{
			return false;
		}
		for (const std::filesystem::path& Rel : AFiles)
		{
			const IoHash HA = IoHash::HashBuffer(ReadFile(A / Rel).Flatten());
			const IoHash HB = IoHash::HashBuffer(ReadFile(B / Rel).Flatten());
			if (HA != HB)
			{
				return false;
			}
		}
		return true;
	}

	struct TestHarness
	{
		TestHarness() : Workers(/*BoostWorkers*/ false, /*SingleThreaded*/ false), Progress(CreateStandardProgress(zen::logging::Default()))
		{
		}

		StorageInstance MakeStorage(const std::filesystem::path& StoragePath)
		{
			StorageInstance SI;
			SI.BuildStorage = CreateFileBuildStorage(StoragePath, StorageStats, /*EnableJsonOutput*/ false);
			return SI;
		}

		std::pair<Oid, Oid> UploadOnce(StorageInstance&				Storage,
									   const std::filesystem::path& SourceFolder,
									   const std::filesystem::path& TempDir)
		{
			const Oid BuildId	  = Oid::NewOid();
			const Oid BuildPartId = Oid::NewOid();

			auto ChunkController = CreateStandardChunkingController(StandardChunkingControllerSettings{});
			auto ChunkCache		 = CreateNullChunkingCache();

			CreateDirectories(TempDir);

			UploadFolderOptions Options{};
			Options.TempDir					  = TempDir;
			Options.FindBlockMaxCount		  = 10000;
			Options.BlockReuseMinPercentLimit = 85;
			Options.AllowMultiparts			  = true;
			Options.CreateBuild				  = true;
			Options.IgnoreExistingBlocks	  = false;
			Options.UploadToZenCache		  = false;
			Options.IsQuiet					  = true;

			const CbObject MetaData;

			UploadFolder(zen::logging::Default(),
						 *Progress,
						 Workers,
						 Storage,
						 AbortFlag,
						 PauseFlag,
						 BuildId,
						 BuildPartId,
						 /*BuildPartName*/ "default"sv,
						 SourceFolder,
						 /*ManifestPath*/ {},
						 MetaData,
						 *ChunkController,
						 *ChunkCache,
						 Options);

			return {BuildId, BuildPartId};
		}

		void DownloadOnce(StorageInstance&			   Storage,
						  const Oid&				   BuildId,
						  const std::filesystem::path& TargetFolder,
						  const std::filesystem::path& ZenFolderPath,
						  const std::filesystem::path& SystemRootDir,
						  const DownloadOptions*	   OverrideOptions = nullptr)
		{
			CreateDirectories(TargetFolder);
			CreateDirectories(ZenFolderPath);
			CreateDirectories(SystemRootDir);

			DownloadOptions Options;
			if (OverrideOptions)
			{
				Options = *OverrideOptions;
			}
			Options.ZenFolderPath = ZenFolderPath;
			Options.SystemRootDir = SystemRootDir;
			Options.IsQuiet		  = true;

			const std::vector<Oid>		   BuildPartIds;
			const std::vector<std::string> BuildPartNames;

			DownloadFolder(zen::logging::Default(),
						   *Progress,
						   Workers,
						   Storage,
						   AbortFlag,
						   PauseFlag,
						   StorageCacheStats,
						   BuildId,
						   BuildPartIds,
						   BuildPartNames,
						   /*DownloadSpecPath*/ {},
						   TargetFolder,
						   Options);
		}

		std::atomic<bool>			  AbortFlag{false};
		std::atomic<bool>			  PauseFlag{false};
		TransferThreadWorkers		  Workers;
		std::unique_ptr<ProgressBase> Progress;
		BuildStorageBase::Statistics  StorageStats;
		BuildStorageCache::Statistics StorageCacheStats;
	};

}  // namespace buildops_test

TEST_SUITE_BEGIN("remotestore.buildoperations");

// Flagship case: one upload + reupload + multiple download variants against
// the same in-process storage. Exercises scavenge, local-chunk copy,
// cached-block reuse, partial-block fetch, and full-block download.
TEST_CASE("buildoperations.roundtrip.full_variations")
{
	using namespace buildops_test;

	ScopedTemporaryDirectory Root;
	TestHarness				 H;

	const std::filesystem::path FolderA		= Root.Path() / "src_a";
	const std::filesystem::path FolderB		= Root.Path() / "src_b";
	const std::filesystem::path StoragePath = Root.Path() / "storage";
	const std::filesystem::path UploadTemp	= Root.Path() / "upload_tmp";
	const std::filesystem::path SystemRoot	= Root.Path() / "sys";

	MakeTestFolder(FolderA, FolderSpec{.Seed = 1});
	MakeTestFolder(FolderB, FolderSpec{.Seed = 1, .DuplicateFileCount = 20});

	CreateDirectories(StoragePath);
	StorageInstance Storage = H.MakeStorage(StoragePath);

	const auto [BuildIdA, PartIdA] = H.UploadOnce(Storage, FolderA, UploadTemp);
	CHECK(BuildIdA != Oid::Zero);

	// Re-upload A: should round-trip without error and still produce a
	// usable build (we verify via a subsequent download).
	const auto [BuildIdA2, PartIdA2] = H.UploadOnce(Storage, FolderA, UploadTemp);
	CHECK(BuildIdA2 != Oid::Zero);

	// Upload B (shares content with A).
	const auto [BuildIdB, PartIdB] = H.UploadOnce(Storage, FolderB, UploadTemp);
	CHECK(BuildIdB != Oid::Zero);

	// Download A into an empty target. Exercises ScheduleFullBlockDownloads +
	// ScheduleLooseChunkWrites.
	{
		const std::filesystem::path Target	 = Root.Path() / "dl_empty";
		const std::filesystem::path ZenState = Target / ".zen";
		H.DownloadOnce(Storage, BuildIdA, Target, ZenState, SystemRoot);
		CHECK(FoldersEquivalent(FolderA, Target));
	}

	// Re-download A after removing some files but keeping the .zen state
	// dir. Exercises ScheduleCachedBlockWrites.
	{
		const std::filesystem::path Target	 = Root.Path() / "dl_cached";
		const std::filesystem::path ZenState = Target / ".zen";
		H.DownloadOnce(Storage, BuildIdA, Target, ZenState, SystemRoot);

		int Deleted = 0;
		for (auto& E : std::filesystem::recursive_directory_iterator(Target))
		{
			if (Deleted >= 5)
				break;
			if (E.is_regular_file())
			{
				const std::filesystem::path Rel = std::filesystem::relative(E.path(), Target);
				if (!Rel.empty() && Rel.begin()->string() == ".zen")
					continue;
				std::error_code Ec;
				std::filesystem::remove(E.path(), Ec);
				if (!Ec)
					++Deleted;
			}
		}
		CHECK(Deleted > 0);

		H.DownloadOnce(Storage, BuildIdA, Target, ZenState, SystemRoot);
		CHECK(FoldersEquivalent(FolderA, Target));
	}

	// Download B into a target pre-seeded with A's content. Exercises
	// ScheduleLocalChunkCopies and ScheduleScavengedSequenceWrites (the two
	// span-capture sites that were fixed).
	{
		const std::filesystem::path Target	 = Root.Path() / "dl_scavenge";
		const std::filesystem::path ZenState = Target / ".zen";
		CopyTreeExcludingZen(FolderA, Target);

		DownloadOptions Opts;
		Opts.EnableTargetFolderScavenging	= true;
		Opts.EnableOtherDownloadsScavenging = true;
		H.DownloadOnce(Storage, BuildIdB, Target, ZenState, SystemRoot, &Opts);

		CHECK(FoldersEquivalent(FolderB, Target));
	}

	// Partial-block mode.
	{
		const std::filesystem::path Target	 = Root.Path() / "dl_partial";
		const std::filesystem::path ZenState = Target / ".zen";

		DownloadOptions Opts;
		Opts.PartialBlockRequestMode = EPartialBlockRequestMode::All;
		H.DownloadOnce(Storage, BuildIdB, Target, ZenState, SystemRoot, &Opts);

		CHECK(FoldersEquivalent(FolderB, Target));
	}
}

// Abort the download before it can do meaningful work. Expected to unwind
// cleanly, not crash or assert.
TEST_CASE("buildoperations.download.abort_midway")
{
	using namespace buildops_test;

	ScopedTemporaryDirectory Root;
	TestHarness				 H;

	const std::filesystem::path Folder		= Root.Path() / "src";
	const std::filesystem::path StoragePath = Root.Path() / "storage";
	const std::filesystem::path UploadTemp	= Root.Path() / "upload_tmp";
	const std::filesystem::path SystemRoot	= Root.Path() / "sys";

	MakeTestFolder(Folder, FolderSpec{.Seed = 42});
	CreateDirectories(StoragePath);
	StorageInstance Storage = H.MakeStorage(StoragePath);

	const auto [BuildId, PartId] = H.UploadOnce(Storage, Folder, UploadTemp);

	const std::filesystem::path Target	 = Root.Path() / "dl_abort";
	const std::filesystem::path ZenState = Target / ".zen";

	H.AbortFlag.store(true);
	CHECK_NOTHROW(H.DownloadOnce(Storage, BuildId, Target, ZenState, SystemRoot));
}

// Empty source folder round-trip: must not crash, must produce an empty
// download target.
TEST_CASE("buildoperations.roundtrip.empty_folder")
{
	using namespace buildops_test;

	ScopedTemporaryDirectory Root;
	TestHarness				 H;

	const std::filesystem::path Folder		= Root.Path() / "empty";
	const std::filesystem::path StoragePath = Root.Path() / "storage";
	const std::filesystem::path UploadTemp	= Root.Path() / "upload_tmp";
	const std::filesystem::path SystemRoot	= Root.Path() / "sys";

	CreateDirectories(Folder);
	CreateDirectories(StoragePath);
	StorageInstance Storage = H.MakeStorage(StoragePath);

	const auto [BuildId, PartId] = H.UploadOnce(Storage, Folder, UploadTemp);

	const std::filesystem::path Target	 = Root.Path() / "dl_empty";
	const std::filesystem::path ZenState = Target / ".zen";
	H.DownloadOnce(Storage, BuildId, Target, ZenState, SystemRoot);

	CHECK(ListRelative(Target).empty());
}

TEST_SUITE_END();

#endif	// ZEN_WITH_TESTS

}  // namespace zen