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

#include "cache_cmd.h"

#include <zencore/compress.h>
#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/thread.h>
#include <zencore/workthreadpool.h>
#include <zenhttp/httpclient.h>
#include <zenhttp/httpcommon.h>
#include <zenutil/cache/cacherequests.h>
#include <zenutil/packageformat.h>

#include <memory>
#include <random>

ZEN_THIRD_PARTY_INCLUDES_START
#include <cpr/cpr.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

namespace {
	IoBuffer CreateRandomBlob(uint64_t Size)
	{
		static uint64_t Seed{0x7CEBF54E45B9F5D1};
		auto			Next = [](uint64_t& seed) {
			   uint64_t z = (seed += UINT64_C(0x9E3779B97F4A7C15));
			   z		  = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9);
			   z		  = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB);
			   return z ^ (z >> 31);
		};

		IoBuffer  Data(Size);
		uint64_t* DataPtr = reinterpret_cast<uint64_t*>(Data.MutableData());
		while (Size > sizeof(uint64_t))
		{
			*DataPtr++ = Next(Seed);
			Size -= sizeof(uint64_t);
		}
		uint64_t ByteNext	 = Next(Seed);
		uint8_t* ByteDataPtr = reinterpret_cast<uint8_t*>(DataPtr);
		while (Size > 0)
		{
			*ByteDataPtr++ = static_cast<uint8_t>(ByteNext & 0xff);
			ByteNext >>= 8;
			Size--;
		}
		return Data;
	};

	CompressedBuffer CompressBlob(IoBuffer&& Buffer)
	{
		return CompressedBuffer::Compress(SharedBuffer(Buffer), OodleCompressor::Mermaid, OodleCompressionLevel::SuperFast);
	}
}  // namespace

DropCommand::DropCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "n", "namespace", "Namespace name", cxxopts::value(m_NamespaceName), "<namespacename>");
	m_Options.add_option("", "b", "bucket", "Bucket name", cxxopts::value(m_BucketName), "<bucketname>");
	m_Options.parse_positional({"namespace", "bucket"});
}

DropCommand::~DropCommand() = default;

int
DropCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw zen::OptionParseException("unable to resolve server specification");
	}

	if (m_NamespaceName.empty())
	{
		throw zen::OptionParseException("Drop command requires a namespace");
	}

	cpr::Session Session;
	if (m_BucketName.empty())
	{
		ZEN_CONSOLE("Dropping cache namespace '{}' from '{}'", m_NamespaceName, m_HostName);
		Session.SetUrl({fmt::format("{}/z$/{}", m_HostName, m_NamespaceName)});
	}
	else
	{
		ZEN_CONSOLE("Dropping cache bucket '{}/{}' from '{}'", m_NamespaceName, m_BucketName, m_HostName);
		Session.SetUrl({fmt::format("{}/z$/{}/{}", m_HostName, m_NamespaceName, m_BucketName)});
	}

	cpr::Response Result = Session.Delete();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("OK: drop succeeded");

		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("Drop failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("Drop failed: {}", Result.error.message);
	}

	return 1;
}

CacheInfoCommand::CacheInfoCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "n", "namespace", "Namespace name", cxxopts::value(m_NamespaceName), "<namespacename>");
	m_Options.add_option("", "b", "bucket", "Bucket name", cxxopts::value(m_BucketName), "<bucketname>");
	m_Options.parse_positional({"namespace", "bucket"});
}

CacheInfoCommand::~CacheInfoCommand() = default;

int
CacheInfoCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw zen::OptionParseException("unable to resolve server specification");
	}

	cpr::Session Session;
	Session.SetHeader(cpr::Header{{"Accept", "application/json"}});
	if (m_HostName.empty())
	{
		ZEN_CONSOLE("Info on cache from '{}'", m_HostName);
		Session.SetUrl({fmt::format("{}/z$", m_HostName)});
	}
	else if (m_BucketName.empty())
	{
		ZEN_CONSOLE("Info on cache namespace '{}' from '{}'", m_NamespaceName, m_HostName);
		Session.SetUrl({fmt::format("{}/z$/{}", m_HostName, m_NamespaceName)});
	}
	else
	{
		ZEN_CONSOLE("Info on cache bucket '{}/{}' from '{}'", m_NamespaceName, m_BucketName, m_HostName);
		Session.SetUrl({fmt::format("{}/z$/{}/{}", m_HostName, m_NamespaceName, m_BucketName)});
	}

	cpr::Response Result = Session.Get();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("{}", Result.text);

		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("Info failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("Info failed: {}", Result.error.message);
	}

	return 1;
}

CacheStatsCommand::CacheStatsCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
}

CacheStatsCommand::~CacheStatsCommand() = default;

int
CacheStatsCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw zen::OptionParseException("unable to resolve server specification");
	}

	cpr::Session Session;
	Session.SetUrl({fmt::format("{}/stats/z$", m_HostName)});
	Session.SetHeader(cpr::Header{{"Accept", "application/json"}});

	cpr::Response Result = Session.Get();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("{}", Result.text);

		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("Info failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("Info failed: {}", Result.error.message);
	}

	return 1;
}

CacheDetailsCommand::CacheDetailsCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "c", "csv", "Info on csv format", cxxopts::value(m_CSV), "<csv>");
	m_Options.add_option("", "d", "details", "Get detailed information about records", cxxopts::value(m_Details), "<details>");
	m_Options.add_option("",
						 "a",
						 "attachmentdetails",
						 "Get detailed information about attachments",
						 cxxopts::value(m_AttachmentDetails),
						 "<attachmentdetails>");
	m_Options.add_option("", "n", "namespace", "Namespace name to get info for", cxxopts::value(m_Namespace), "<namespace>");
	m_Options.add_option("", "b", "bucket", "Filter on bucket name", cxxopts::value(m_Bucket), "<bucket>");
	m_Options.add_option("", "v", "valuekey", "Filter on value key hash string", cxxopts::value(m_ValueKey), "<valuekey>");
}

CacheDetailsCommand::~CacheDetailsCommand() = default;

int
CacheDetailsCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw zen::OptionParseException("unable to resolve server specification");
	}

	cpr::Session	Session;
	cpr::Parameters Parameters;
	if (m_Details)
	{
		Parameters.Add({"details", "true"});
	}
	if (m_AttachmentDetails)
	{
		Parameters.Add({"attachmentdetails", "true"});
	}
	if (m_CSV)
	{
		Parameters.Add({"csv", "true"});
	}
	else
	{
		Session.SetHeader(cpr::Header{{"Accept", "application/json"}});
	}

	if (!m_ValueKey.empty())
	{
		if (m_Namespace.empty() || m_Bucket.empty())
		{
			ZEN_ERROR("Provide namespace and bucket name");
			ZEN_CONSOLE("{}", m_Options.help({""}).c_str());
			return 1;
		}
		Session.SetUrl({fmt::format("{}/z$/details$/{}/{}/{}", m_HostName, m_Namespace, m_Bucket, m_ValueKey)});
	}
	else if (!m_Bucket.empty())
	{
		if (m_Namespace.empty())
		{
			ZEN_ERROR("Provide namespace name");
			ZEN_CONSOLE("{}", m_Options.help({""}).c_str());
			return 1;
		}
		Session.SetUrl({fmt::format("{}/z$/details$/{}/{}", m_HostName, m_Namespace, m_Bucket)});
	}
	else if (!m_Namespace.empty())
	{
		Session.SetUrl({fmt::format("{}/z$/details$/{}", m_HostName, m_Namespace)});
	}
	else
	{
		Session.SetUrl({fmt::format("{}/z$/details$", m_HostName)});
	}
	Session.SetParameters(Parameters);

	cpr::Response Result = Session.Get();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("{}", Result.text);

		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("Info failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("Info failed: {}", Result.error.message);
	}

	return 1;
}

CacheGenerateCommand::CacheGenerateCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options
		.add_option("", "n", "namespace", "Namespace to generate cache values/records for", cxxopts::value(m_Namespace), "<namespace>");
	m_Options.add_option("", "b", "bucket", "Bucket name to generate cache values/records for", cxxopts::value(m_Bucket), "<bucket>");
	m_Options.add_option("", "", "count", "Number of cache values/records to generate", cxxopts::value(m_Count), "<count>");
	m_Options.add_option("", "", "min-size", "Minimum size of cache value/attachments", cxxopts::value(m_MinSize), "<min>");
	m_Options.add_option("", "", "max-size", "Maximum size of cache value/attachments", cxxopts::value(m_MaxSize), "<max>");
	m_Options.add_option("",
						 "",
						 "min-attachments",
						 "Minimum number of attachments when creating record based values",
						 cxxopts::value(m_MinAttachmentCount),
						 "<minattachments>");
	m_Options.add_option("",
						 "",
						 "max-attachments",
						 "Minimum number of attachments when creating record based values, 0 to only create cache values",
						 cxxopts::value(m_MaxAttachmentCount),
						 "<maxattachments>");
	m_Options.parse_positional({"namespace", "bucket", "count"});
	m_Options.positional_help("namespace bucket count");
}

CacheGenerateCommand::~CacheGenerateCommand() = default;

int
CacheGenerateCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw zen::OptionParseException("unable to resolve server specification");
	}

	if (m_MaxSize == 0 && m_MinSize == 0)
	{
		m_MinSize = 17;
		if (m_MaxAttachmentCount == 0)
		{
			// For cache values this max size will result in about 0.5% of values being saved as loose file in a cache bucket
			m_MaxSize = 65u * 1024u;
		}
		else
		{
			// For cache records this max size will result in about 0.5% of attachments begin saved as loose files in cas
			m_MaxSize = 768u * 1024u;
		}
	}

	std::vector<std::uniform_int_distribution<uint64_t>> Variations;
	std::vector<size_t>									 SizeRanges;
	SizeRanges.push_back(m_MinSize);
	Variations.push_back(std::uniform_int_distribution<uint64_t>(0, m_MinSize - 1));
	while (SizeRanges.back() < m_MaxSize)
	{
		SizeRanges.push_back(SizeRanges.back() * 2);
		Variations.push_back(std::uniform_int_distribution<uint64_t>(0, SizeRanges.back() - 1));
	}
	if (SizeRanges.back() > m_MaxSize)
	{
		SizeRanges.back() = m_MaxSize;
		Variations.push_back(std::uniform_int_distribution<uint64_t>(0, m_MaxSize - 1));
	}

	std::random_device						RandomDevice;
	std::mt19937							Generator(RandomDevice());
	std::uniform_int_distribution<uint64_t> SizeRangeDistribution(0, SizeRanges.size() - 1);

	std::vector<uint64_t> Sizes;
	Sizes.reserve(m_Count);
	for (uint64_t n = 0; n != m_Count; ++n)
	{
		uint64_t Range	   = SizeRangeDistribution(Generator);
		uint64_t Size	   = SizeRanges[Range];
		uint64_t Variation = Variations[Range](Generator);
		Sizes.push_back(Size + Variation);
	}

	std::uniform_int_distribution<uint64_t> KeyDistribution;

	HttpClient Http(m_HostName);

	auto GeneratePutCacheValueRequest([this, &KeyDistribution, &Generator](std::span<std::uint64_t> BatchSizes, uint64_t RequestIndex) {
		cacherequests::PutCacheValuesRequest Request({.AcceptMagic = kCbPkgMagic, .Namespace = m_Namespace});
		for (std::uint64_t ValueSize : BatchSizes)
		{
			uint64_t	KeyBase	  = KeyDistribution(Generator);
			std::string KeyString = fmt::format("{}-{}-{}", RequestIndex, KeyBase, ValueSize);
			IoHash		ValueKey  = IoHash::HashBuffer(KeyString.c_str(), KeyString.length());

			Request.Requests.emplace_back(cacherequests::PutCacheValueRequest{.Key	= {.Bucket = m_Bucket, .Hash = ValueKey},
																			  .Body = CompressBlob(CreateRandomBlob(ValueSize))});
		}
		return Request;
	});

	auto GeneratePutCacheRecordRequest([this, &KeyDistribution, &Generator](std::span<std::uint64_t> BatchSizes, uint64_t RequestIndex) {
		cacherequests::PutCacheRecordsRequest Request({.AcceptMagic = kCbPkgMagic, .Namespace = m_Namespace});
		uint64_t							  KeyBase		  = KeyDistribution(Generator);
		std::string							  RecordKeyString = fmt::format("{}-{}-{}", RequestIndex, KeyBase, BatchSizes.size());
		IoHash								  RecordKey		  = IoHash::HashBuffer(RecordKeyString.c_str(), RecordKeyString.length());

		Request.Requests.emplace_back(cacherequests::PutCacheRecordRequest{.Key = {.Bucket = m_Bucket, .Hash = RecordKey}});
		for (std::uint64_t ValueSize : BatchSizes)
		{
			Request.Requests.back().Values.push_back({.Id = Oid::NewOid(), .Body = CompressBlob(CreateRandomBlob(ValueSize))});
		}
		return Request;
	});

	WorkerThreadPool WorkerPool(gsl::narrow<int>(Max((std::thread::hardware_concurrency() / 2u), 2u)));
	Latch			 WorkLatch(1);

	std::uniform_int_distribution<uint32_t> SizeCountDistribution(m_MaxAttachmentCount > 0 ? 0 : 1,
																  m_MaxAttachmentCount > 0 ? m_MaxAttachmentCount : 8);

	std::size_t Offset		 = 0;
	uint64_t	RequestIndex = 0;
	while (Offset < Sizes.size())
	{
		size_t				SizeCount  = SizeCountDistribution(Generator);
		std::span<uint64_t> BatchSizes = std::span<uint64_t>(Sizes).subspan(Offset, Min(Max(SizeCount, 1u), Sizes.size() - Offset));

		WorkLatch.AddCount(1);
		WorkerPool.ScheduleWork([&, BatchSizes, RequestIndex]() {
			auto	  _ = MakeGuard([&WorkLatch]() { WorkLatch.CountDown(); });
			CbPackage Package;
			if (m_MaxAttachmentCount > 0 && SizeCount > 0)
			{
				auto Request = GeneratePutCacheRecordRequest(BatchSizes, RequestIndex);
				ZEN_ASSERT(Request.Format(Package));
			}
			else
			{
				auto Request = GeneratePutCacheValueRequest(BatchSizes, RequestIndex);
				ZEN_ASSERT(Request.Format(Package));
			}

			if (HttpClient::Response Response = Http.Post("/z$/$rpc", Package, HttpClient::Accept(ZenContentType::kCbPackage)); !Response)
			{
				ZEN_CONSOLE("{}", Response.ErrorMessage(fmt::format("{}: ", RequestIndex)));
			}
		});
		Offset += BatchSizes.size();
		RequestIndex++;
	}

	WorkLatch.CountDown();
	while (!WorkLatch.Wait(1000))
	{
		ZEN_INFO("Creating data, {} requests remaining", WorkLatch.Remaining());
	}

	return 0;
}

}  // namespace zen