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

#include <zencore/compactbinary.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/compress.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/stream.h>

#if ZEN_WITH_TESTS
#	define ZEN_TEST_WITH_RUNNER 1
#	include <zencore/testing.h>
#endif

#include <fmt/format.h>

#include <stdio.h>
#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <string>
#include <system_error>
#include <thread>

using namespace std::literals;
using namespace zen;

#if !defined(_MSC_VER)
#	define _strnicmp strncasecmp  // TEMPORARY WORKAROUND - should not be using this
#endif

// Some basic functions to implement some test "compute" functions

std::string
Rot13Function(std::string_view InputString)
{
	std::string OutputString{InputString};

	std::transform(OutputString.begin(),
				   OutputString.end(),
				   OutputString.begin(),
				   [](std::string::value_type c) -> std::string::value_type {
					   if (c >= 'a' && c <= 'z')
					   {
						   return 'a' + (c - 'a' + 13) % 26;
					   }
					   else if (c >= 'A' && c <= 'Z')
					   {
						   return 'A' + (c - 'A' + 13) % 26;
					   }
					   else
					   {
						   return c;
					   }
				   });

	return OutputString;
}

std::string
ReverseFunction(std::string_view InputString)
{
	std::string OutputString{InputString};
	std::reverse(OutputString.begin(), OutputString.end());
	return OutputString;
}

std::string
IdentityFunction(std::string_view InputString)
{
	return std::string{InputString};
}

std::string
NullFunction(std::string_view)
{
	return {};
}

zen::CbObject
DescribeFunctions()
{
	CbObjectWriter Versions;
	Versions << "BuildSystemVersion" << Guid::FromString("17fe280d-ccd8-4be8-a9d1-89c944a70969"sv);

	Versions.BeginArray("Functions"sv);
	Versions.BeginObject();
	Versions << "Name"sv
			 << "Null"sv;
	Versions << "Version"sv << Guid::FromString("00000000-0000-0000-0000-000000000000"sv);
	Versions.EndObject();
	Versions.BeginObject();
	Versions << "Name"sv
			 << "Identity"sv;
	Versions << "Version"sv << Guid::FromString("11111111-1111-1111-1111-111111111111"sv);
	Versions.EndObject();
	Versions.BeginObject();
	Versions << "Name"sv
			 << "Rot13"sv;
	Versions << "Version"sv << Guid::FromString("13131313-1313-1313-1313-131313131313"sv);
	Versions.EndObject();
	Versions.BeginObject();
	Versions << "Name"sv
			 << "Reverse"sv;
	Versions << "Version"sv << Guid::FromString("31313131-3131-3131-3131-313131313131"sv);
	Versions.EndObject();
	Versions.EndArray();

	return Versions.Save();
}

struct ContentResolver
{
	std::filesystem::path InputsRoot;

	CompressedBuffer ResolveChunk(IoHash Hash, uint64_t ExpectedSize)
	{
		std::filesystem::path ChunkPath	  = InputsRoot / Hash.ToHexString();
		IoBuffer			  ChunkBuffer = IoBufferBuilder::MakeFromFile(ChunkPath);

		IoHash			 RawHash;
		uint64_t		 RawSize	  = 0;
		CompressedBuffer AsCompressed = CompressedBuffer::FromCompressed(SharedBuffer(ChunkBuffer), RawHash, RawSize);

		if (RawSize != ExpectedSize)
		{
			throw std::runtime_error(
				fmt::format("chunk size mismatch - expected {}, got {} for '{}'", ExpectedSize, ChunkBuffer.Size(), ChunkPath));
		}
		if (RawHash != Hash)
		{
			throw std::runtime_error(fmt::format("chunk hash mismatch - expected {}, got {} for '{}'", Hash, RawHash, ChunkPath));
		}

		return AsCompressed;
	}
};

zen::CbPackage
ExecuteFunction(CbObject Action, ContentResolver ChunkResolver)
{
	auto Apply = [&](auto Func) {
		zen::CbPackage Result;
		auto		   Source = Action["Inputs"sv].AsObjectView()["Source"sv].AsObjectView();

		IoHash	 InputRawHash = Source["RawHash"sv].AsHash();
		uint64_t InputRawSize = Source["RawSize"sv].AsUInt64();

		zen::CompressedBuffer InputData = ChunkResolver.ResolveChunk(InputRawHash, InputRawSize);
		SharedBuffer		  Input		= InputData.Decompress();

		std::string			  Output = Func(std::string_view(static_cast<const char*>(Input.GetData()), Input.GetSize()));
		zen::CompressedBuffer OutputData =
			zen::CompressedBuffer::Compress(SharedBuffer::MakeView(Output), OodleCompressor::Selkie, OodleCompressionLevel::HyperFast4);
		IoHash OutputRawHash = OutputData.DecodeRawHash();

		CbAttachment OutputAttachment(std::move(OutputData), OutputRawHash);

		CbObjectWriter Cbo;
		Cbo.BeginArray("Values"sv);
		Cbo.BeginObject();
		Cbo << "Id" << Oid{1, 2, 3};
		Cbo.AddAttachment("RawHash", OutputAttachment);
		Cbo << "RawSize" << Output.size();
		Cbo.EndObject();
		Cbo.EndArray();

		Result.SetObject(Cbo.Save());
		Result.AddAttachment(std::move(OutputAttachment));
		return Result;
	};

	std::string_view Function = Action["Function"sv].AsString();

	if (Function == "Rot13"sv)
	{
		return Apply(Rot13Function);
	}
	else if (Function == "Reverse"sv)
	{
		return Apply(ReverseFunction);
	}
	else if (Function == "Identity"sv)
	{
		return Apply(IdentityFunction);
	}
	else if (Function == "Null"sv)
	{
		return Apply(NullFunction);
	}
	else
	{
		return {};
	}
}

/* This implements a minimal application to help testing of process launch-related
   functionality

   It also mimics the DDC2 worker command line interface, so it may be used to
   exercise compute infrastructure.
 */

int
main(int argc, char* argv[])
{
	int ExitCode = 0;

	try
	{
		std::filesystem::path			   BasePath	   = std::filesystem::current_path();
		std::filesystem::path			   InputPath   = std::filesystem::current_path() / "Inputs";
		std::filesystem::path			   OutputPath  = std::filesystem::current_path() / "Outputs";
		std::filesystem::path			   VersionPath = std::filesystem::current_path() / "Versions";
		std::vector<std::filesystem::path> ActionPaths;

		/*
					GetSwitchValues(TEXT("-B="), ActionPathPatterns);
					GetSwitchValues(TEXT("-Build="), ActionPathPatterns);

					GetSwitchValues(TEXT("-I="), InputDirectoryPaths);
					GetSwitchValues(TEXT("-Input="), InputDirectoryPaths);

					GetSwitchValues(TEXT("-O="), OutputDirectoryPaths);
					GetSwitchValues(TEXT("-Output="), OutputDirectoryPaths);

					GetSwitchValues(TEXT("-V="), VersionPaths);
					GetSwitchValues(TEXT("-Version="), VersionPaths);
		 */

		auto SplitArg = [](const char* Arg) -> std::string_view {
			std::string_view ArgView{Arg};
			if (auto SplitPos = ArgView.find_first_of('='); SplitPos != std::string_view::npos)
			{
				return ArgView.substr(SplitPos + 1);
			}
			else
			{
				return {};
			}
		};

		auto ParseIntArg = [](std::string_view Arg) -> int {
			int		   Rv	  = 0;
			const auto Result = std::from_chars(Arg.data(), Arg.data() + Arg.size(), Rv);

			if (Result.ec != std::errc{})
			{
				throw std::invalid_argument(fmt::format("bad argument (not an integer): {}", Arg).c_str());
			}

			return Rv;
		};

		for (int i = 1; i < argc; ++i)
		{
			std::string_view Arg = argv[i];

			if (Arg.compare(0, 1, "-"))
			{
				continue;
			}

			if (std::strncmp(argv[i], "-t=", 3) == 0)
			{
				const int SleepTime = std::atoi(argv[i] + 3);

				printf("[zentest] sleeping for %ds...\n", SleepTime);

				std::this_thread::sleep_for(SleepTime * 1s);
			}
			else if (std::strncmp(argv[i], "-f=", 3) == 0)
			{
				// Force a "failure" process exit code to return to the invoker

				// This may throw for invalid arguments, which makes this useful for
				// testing exception handling
				std::string_view ErrorArg = SplitArg(argv[i]);
				ExitCode				  = ParseIntArg(ErrorArg);
			}
			else if ((_strnicmp(argv[i], "-input=", 7) == 0) || (_strnicmp(argv[i], "-i=", 3) == 0))
			{
				/* mimic DDC2

					GetSwitchValues(TEXT("-I="), InputDirectoryPaths);
					GetSwitchValues(TEXT("-Input="), InputDirectoryPaths);
				*/

				std::string_view InputArg = SplitArg(argv[i]);
				InputPath				  = InputArg;
			}
			else if ((_strnicmp(argv[i], "-output=", 8) == 0) || (_strnicmp(argv[i], "-o=", 3) == 0))
			{
				/* mimic DDC2 handling of where files storing output chunk files are directed

					GetSwitchValues(TEXT("-O="), OutputDirectoryPaths);
					GetSwitchValues(TEXT("-Output="), OutputDirectoryPaths);
				*/

				std::string_view OutputArg = SplitArg(argv[i]);
				OutputPath				   = OutputArg;
			}
			else if ((_strnicmp(argv[i], "-version=", 8) == 0) || (_strnicmp(argv[i], "-v=", 3) == 0))
			{
				/* mimic DDC2

					GetSwitchValues(TEXT("-V="), VersionPaths);
					GetSwitchValues(TEXT("-Version="), VersionPaths);
				*/

				std::string_view VersionArg = SplitArg(argv[i]);
				VersionPath					= VersionArg;
			}
			else if ((_strnicmp(argv[i], "-build=", 7) == 0) || (_strnicmp(argv[i], "-b=", 3) == 0))
			{
				/* mimic DDC2

					GetSwitchValues(TEXT("-B="), ActionPathPatterns);
					GetSwitchValues(TEXT("-Build="), ActionPathPatterns);
				*/

				std::string_view	  BuildActionArg = SplitArg(argv[i]);
				std::filesystem::path ActionPath{BuildActionArg};
				ActionPaths.push_back(ActionPath);

				ExitCode = 0;
			}
		}

		// Emit version information

		if (!VersionPath.empty())
		{
			CbObjectWriter Version;

			Version << "BuildSystemVersion" << Guid::FromString("17fe280d-ccd8-4be8-a9d1-89c944a70969"sv);

			Version.BeginArray("Functions");

			Version.BeginObject();
			Version << "Name"
					<< "Rot13"
					<< "Version" << Guid::FromString("13131313-1313-1313-1313-131313131313"sv);
			Version.EndObject();

			Version.BeginObject();
			Version << "Name"
					<< "Reverse"
					<< "Version" << Guid::FromString("98765432-1000-0000-0000-000000000000"sv);
			Version.EndObject();

			Version.BeginObject();
			Version << "Name"
					<< "Identity"
					<< "Version" << Guid::FromString("11111111-1111-1111-1111-111111111111"sv);
			Version.EndObject();

			Version.BeginObject();
			Version << "Name"
					<< "Null"
					<< "Version" << Guid::FromString("00000000-0000-0000-0000-000000000000"sv);
			Version.EndObject();

			Version.EndArray();
			CbObject VersionObject = Version.Save();

			BinaryWriter Writer;
			zen::SaveCompactBinary(Writer, VersionObject);
			zen::WriteFile(VersionPath, IoBufferBuilder::MakeFromMemory(Writer.GetView()));
		}

		// Evaluate actions

		ContentResolver Resolver;
		Resolver.InputsRoot = InputPath;

		for (std::filesystem::path ActionPath : ActionPaths)
		{
			IoBuffer  ActionDescBuffer = ReadFile(ActionPath).Flatten();
			CbObject  ActionDesc	   = LoadCompactBinaryObject(ActionDescBuffer);
			CbPackage Result		   = ExecuteFunction(ActionDesc, Resolver);
			CbObject  ResultObject	   = Result.GetObject();

			BinaryWriter Writer;
			zen::SaveCompactBinary(Writer, ResultObject);
			zen::WriteFile(ActionPath.replace_extension(".output"), IoBufferBuilder::MakeFromMemory(Writer.GetView()));

			// Also marshal outputs

			for (const auto& Attachment : Result.GetAttachments())
			{
				const CompositeBuffer& AttachmentBuffer = Attachment.AsCompressedBinary().GetCompressed();
				zen::WriteFile(OutputPath / Attachment.GetHash().ToHexString(), AttachmentBuffer.Flatten().AsIoBuffer());
			}
		}
	}
	catch (std::exception& Ex)
	{
		printf("[zentest] exception caught in main: '%s'\n", Ex.what());

		ExitCode = 99;
	}

	printf("[zentest] exiting with exit code: %d\n", ExitCode);

	return ExitCode;
}