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

#pragma once

#if ZEN_WITH_TESTS

#	include <zencore/compactbinarybuilder.h>
#	include <zencore/compactbinarypackage.h>
#	include <zencore/iobuffer.h>
#	include <zencore/stream.h>
#	include <zencore/string.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#	include <zenhttp/httpcommon.h>
#	include <zenutil/zenserverprocess.h>

#	include <functional>

namespace zen::tests {

extern zen::ZenServerEnvironment TestEnv;

inline IoBuffer
MakeCbObjectPayload(std::function<void(CbObjectWriter& Writer)> WriteCB)
{
	CbObjectWriter Writer;
	WriteCB(Writer);
	IoBuffer Payload = Writer.Save().GetBuffer().AsIoBuffer();
	Payload.SetContentType(ZenContentType::kCbObject);
	return Payload;
};

inline IoBuffer
SerializeToBuffer(const zen::CbPackage& Package)
{
	BinaryWriter MemStream;

	Package.Save(MemStream);

	IoBuffer Buffer = zen::IoBuffer(zen::IoBuffer::Clone, MemStream.Data(), MemStream.Size());
	Buffer.SetContentType(HttpContentType::kCbPackage);
	return Buffer;
};

namespace utils {

	struct ZenConfig
	{
		std::filesystem::path DataDir;
		uint16_t			  Port;
		std::string			  BaseUri;
		std::string			  Args;

		static ZenConfig New(std::string Args = "")
		{
			return ZenConfig{.DataDir = TestEnv.CreateNewTestDir(), .Port = TestEnv.GetNewPortNumber(), .Args = std::move(Args)};
		}

		static ZenConfig New(uint16_t Port, std::string Args = "")
		{
			return ZenConfig{.DataDir = TestEnv.CreateNewTestDir(), .Port = Port, .Args = std::move(Args)};
		}

		static ZenConfig NewWithUpstream(uint16_t Port, uint16_t UpstreamPort, std::string Args = "")
		{
			return New(Port,
					   fmt::format("{}{}--debug --upstream-thread-count=0 --upstream-zen-url=http://localhost:{}",
								   Args,
								   Args.length() > 0 ? " " : "",
								   UpstreamPort));
		}

		static ZenConfig NewWithThreadedUpstreams(uint16_t NewPort, std::span<uint16_t> UpstreamPorts, bool Debug)
		{
			std::string Args = Debug ? "--debug" : "";
			for (uint16_t Port : UpstreamPorts)
			{
				Args = fmt::format("{}{}--upstream-zen-url=http://localhost:{}", Args, Args.length() > 0 ? " " : "", Port);
			}
			return New(NewPort, Args);
		}

		void Spawn(ZenServerInstance& Inst)
		{
			Inst.SetDataDir(DataDir);
			Inst.SpawnServer(Port, Args);
			const uint16_t InstancePort = Inst.WaitUntilReady();
			CHECK_MESSAGE(InstancePort != 0, Inst.GetLogOutput());

			if (Port != InstancePort)
				ZEN_DEBUG("relocation detected from {} to {}", Port, InstancePort);

			Port	= InstancePort;
			BaseUri = fmt::format("http://localhost:{}/z$", Port);
		}
	};

	inline void SpawnServer(ZenServerInstance& Server, ZenConfig& Cfg) { Cfg.Spawn(Server); }

	inline CompressedBuffer CreateSemiRandomBlob(size_t				   AttachmentSize,
												 OodleCompressionLevel CompressionLevel = OodleCompressionLevel::VeryFast)
	{
		// Convoluted way to get a compressed buffer whose result it large enough to be a separate file
		// but also does actually compress
		const size_t			  PartCount = (AttachmentSize / (1u * 1024u * 64)) + 1;
		const size_t			  PartSize	= AttachmentSize / PartCount;
		auto					  Part		= SharedBuffer(CreateRandomBlob(PartSize));
		std::vector<SharedBuffer> Parts(PartCount, Part);
		size_t					  RemainPartSize = AttachmentSize - (PartSize * PartCount);
		if (RemainPartSize > 0)
		{
			Parts.push_back(SharedBuffer(CreateRandomBlob(RemainPartSize)));
		}
		CompressedBuffer Value = CompressedBuffer::Compress(CompositeBuffer(std::move(Parts)), OodleCompressor::Mermaid, CompressionLevel);
		return Value;
	};

	inline std::vector<std::pair<Oid, CompressedBuffer>> CreateAttachments(const std::span<const size_t>& Sizes)
	{
		std::vector<std::pair<Oid, CompressedBuffer>> Result;
		Result.reserve(Sizes.size());
		for (size_t Size : Sizes)
		{
			CompressedBuffer Compressed = CompressedBuffer::Compress(SharedBuffer(CreateRandomBlob(Size)));
			Result.emplace_back(std::pair<Oid, CompressedBuffer>(Oid::NewOid(), Compressed));
		}
		return Result;
	}

	inline std::vector<std::pair<Oid, CompressedBuffer>> CreateSemiRandomAttachments(const std::span<const size_t>& Sizes)
	{
		std::vector<std::pair<Oid, CompressedBuffer>> Result;
		Result.reserve(Sizes.size());
		for (size_t Size : Sizes)
		{
			CompressedBuffer Compressed =
				CreateSemiRandomBlob(Size, Size > 1024u * 1024u ? OodleCompressionLevel::None : OodleCompressionLevel::VeryFast);
			Result.emplace_back(std::pair<Oid, CompressedBuffer>(Oid::NewOid(), Compressed));
		}
		return Result;
	}

}  // namespace utils

class ZenServerTestHelper
{
public:
	ZenServerTestHelper(std::string_view HelperId, int ServerCount) : m_HelperId{HelperId}, m_ServerCount{ServerCount} {}
	~ZenServerTestHelper() {}

	void SpawnServers(std::string_view AdditionalServerArgs = std::string_view())
	{
		SpawnServers([](ZenServerInstance&) {}, AdditionalServerArgs);
	}

	void SpawnServers(auto&& Callback, std::string_view AdditionalServerArgs)
	{
		ZEN_INFO("{}: spawning {} server instances", m_HelperId, m_ServerCount);

		m_Instances.resize(m_ServerCount);

		for (int i = 0; i < m_ServerCount; ++i)
		{
			auto& Instance = m_Instances[i];
			Instance	   = std::make_unique<ZenServerInstance>(TestEnv);
			Instance->SetDataDir(TestEnv.CreateNewTestDir());
		}

		for (int i = 0; i < m_ServerCount; ++i)
		{
			auto& Instance = m_Instances[i];
			Callback(*Instance);
		}

		for (int i = 0; i < m_ServerCount; ++i)
		{
			auto& Instance = m_Instances[i];
			Instance->SpawnServer(TestEnv.GetNewPortNumber(), AdditionalServerArgs);
		}

		for (int i = 0; i < m_ServerCount; ++i)
		{
			auto&	 Instance	= m_Instances[i];
			uint16_t PortNumber = Instance->WaitUntilReady();
			CHECK_MESSAGE(PortNumber != 0, Instance->GetLogOutput());
		}
	}

	ZenServerInstance& GetInstance(int Index) { return *m_Instances[Index]; }

private:
	std::string										m_HelperId;
	int												m_ServerCount = 0;
	std::vector<std::unique_ptr<ZenServerInstance>> m_Instances;
};

inline std::string
OidAsString(const Oid& Id)
{
	StringBuilder<25> OidStringBuilder;
	Id.ToString(OidStringBuilder);
	return OidStringBuilder.ToString();
}

}  // namespace zen::tests

#endif