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

#include "importproject.h"
#include "exportproject.h"

#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/logging.h>
#include <zencore/stream.h>
#include <zenhttp/httpcommon.h>
#include <zenhttp/httpshared.h>
#include <zenutil/basicfile.h>

#include <memory>

ZEN_THIRD_PARTY_INCLUDES_START
#include <cpr/cpr.h>
ZEN_THIRD_PARTY_INCLUDES_END

ImportProjectCommand::ImportProjectCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value("http://localhost:1337"), "<hosturl>");
	m_Options.add_option("", "s", "source", "Source path", cxxopts::value(m_SourcePath), "<sourcepath>");
	m_Options.add_option("", "p", "project", "Project name", cxxopts::value(m_ProjectName), "<projectname>");
	m_Options.add_option("", "o", "oplog", "Oplog name", cxxopts::value(m_OplogNames), "<oplog>");
}

ImportProjectCommand::~ImportProjectCommand() = default;

int
ImportProjectCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	using namespace std::literals;

	ZEN_UNUSED(GlobalOptions);

	m_Options.parse_positional({"source", "project", "oplog"});
	m_Options.parse(argc, argv);

	if (m_ProjectName.empty())
	{
		ZEN_ERROR("Project name must be given");
		return 1;
	}

	if (m_SourcePath.empty())
	{
		ZEN_ERROR("Source path must be given");
		return 1;
	}

	if (!std::filesystem::is_directory(m_SourcePath))
	{
		ZEN_ERROR("Source path '{}' is not a directory", m_SourcePath);
		return 1;
	}

	const std::string UrlBase = fmt::format("{}/prj", m_HostName);

	cpr::Session Session;

	{
		ZEN_INFO("Requesting project '{}' from '{}'", m_ProjectName, m_HostName);

		zen::BasicFile ProjectStore;
		ProjectStore.Open(ExportProjectCommand::GetProjectPath(m_SourcePath, m_ProjectName), zen::BasicFile::Mode::kRead);
		zen::IoBuffer Payload = ProjectStore.ReadAll();

		std::string ProjectRequest = fmt::format("{}/{}", UrlBase, m_ProjectName);
		Session.SetUrl({ProjectRequest});
		Session.SetBody(cpr::Body{(const char*)Payload.GetData(), Payload.GetSize()});
		cpr::Response Response = Session.Post();
		if (!ExportProjectCommand::IsSuccess(Response, ProjectRequest))
		{
			return 1;
		}
		if (m_OplogNames.empty())
		{
			zen::DirectoryContent Content;
			zen::GetDirectoryContent(m_SourcePath, zen::DirectoryContent::IncludeFilesFlag, Content);
			for (auto& File : Content.Files)
			{
				if (File.extension() == ".ops")
				{
					m_OplogNames.push_back(File.stem().string());
				}
			}
		}
	}

	zen::IoBuffer ChunkStoreIndex = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksIndexPath(m_SourcePath));
	zen::IoBuffer ChunkStoreHeaderMem(ChunkStoreIndex, 0, sizeof(ExportProjectCommand::ChunksHeader));
	const ExportProjectCommand::ChunksHeader* ChunksHeader =
		reinterpret_cast<const ExportProjectCommand::ChunksHeader*>(ChunkStoreHeaderMem.GetView().GetData());

	if (ChunksHeader->Magic != ExportProjectCommand::ChunksHeader::kMagic)
	{
		ZEN_ERROR("Invalid chunk index header");
		return 1;
	}
	const size_t BlockSize = 1ull << ChunksHeader->BlockSizeShift;

	zen::IoBuffer							ChunkStoreEntriesMem(ChunkStoreIndex,
										 sizeof(ExportProjectCommand::ChunksHeader),
										 sizeof(ExportProjectCommand::ChunkEntry) * ChunksHeader->ChunkCount);
	const ExportProjectCommand::ChunkEntry* ChunkEntries =
		reinterpret_cast<const ExportProjectCommand::ChunkEntry*>(ChunkStoreEntriesMem.GetView().GetData());

	for (const std::string& OplogName : m_OplogNames)
	{
		ZEN_INFO("Importing oplog '{}/{}' from '{}' to '{}'", m_ProjectName, OplogName, m_SourcePath, m_HostName);
		std::string GetOplogRequest = fmt::format("{}/{}/oplog/{}", UrlBase, m_ProjectName, OplogName);
		Session.SetUrl(GetOplogRequest);
		cpr::Response OplogResponse = Session.Get();
		if (OplogResponse.status_code == static_cast<long>(zen::HttpResponseCode::NotFound))
		{
			OplogResponse = Session.Post();
			if (!ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest))
			{
				return 1;
			}
			ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest);
			OplogResponse = Session.Get();
			if (!ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest))
			{
				return 1;
			}
		}

		zen::BasicFile OpStore;
		OpStore.Open(ExportProjectCommand::GetOplogPath(m_SourcePath, OplogName), zen::BasicFile::Mode::kRead);
		ExportProjectCommand::OplogHeader OplogHeader;
		OpStore.Read(&OplogHeader, sizeof(ExportProjectCommand::OplogHeader), 0);
		if (OplogHeader.Magic != ExportProjectCommand::OplogHeader::kMagic ||
			OplogHeader.HeaderSize != sizeof(ExportProjectCommand::OplogHeader))
		{
			ZEN_ERROR("Invalid oplog header");
			return 1;
		}
		zen::IoHash									  Checksum = OplogHeader.Checksum;
		std::vector<ExportProjectCommand::OplogEntry> OpEntries;
		OpEntries.resize(OplogHeader.OpCount);
		OpStore.Read(OpEntries.data(),
					 sizeof(ExportProjectCommand::OplogEntry) * OplogHeader.OpCount,
					 sizeof(ExportProjectCommand::OplogHeader));

		ZEN_INFO("Constructing oplog with {} ops with checksum '{}' for '{}/{}'",
				 OpEntries.size(),
				 OplogHeader.Checksum,
				 m_ProjectName,
				 OplogName);
		zen::IoHashStream	Hasher;
		zen::CbObjectWriter Request;
		Request.BeginArray("entries"sv);
		for (auto& OpEntry : OpEntries)
		{
			zen::IoBuffer CoreData(OpEntry.OpLength);
			OpStore.Read(CoreData.MutableData(), OpEntry.OpLength, OpEntry.Offset);
			Hasher.Append(CoreData.GetView());
			zen::SharedBuffer SharedCoreData(CoreData);

			zen::CbObject Op(SharedCoreData);
			Request << Op;
		}
		Request.EndArray();
		zen::IoHash CalculatedChecksum = Hasher.GetHash();
		if (CalculatedChecksum != Checksum)
		{
			ZEN_ERROR("Checksum for oplog does not match. Expected '{}' but got '{}'", Checksum, CalculatedChecksum);
			return 1;
		}
		Request.AddHash("checksum"sv, Checksum);

		zen::CbPackage RequestPackage;
		RequestPackage.SetObject(Request.Save());

		ZEN_INFO("Assembling {} attachments", ChunksHeader->ChunkCount);
		std::vector<zen::CbAttachment> Attachments;
		Attachments.reserve(ChunksHeader->ChunkCount);
		uint32_t	  ReadBlockIndex = 0;
		zen::IoBuffer BlockStore = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksPath(m_SourcePath, ReadBlockIndex));
		for (uint64_t ChunkIndex = 0; ChunkIndex < ChunksHeader->ChunkCount; ++ChunkIndex)
		{
			const ExportProjectCommand::ChunkEntry& ChunkEntry = ChunkEntries[ChunkIndex];
			if (ChunkEntry.Offset == ~0ull)
			{
				zen::IoBuffer ChunkBuffer =
					zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetLargeChunkPath(m_SourcePath, ChunkEntry.ChunkHash));
				zen::CompressedBuffer Chunk = zen::CompressedBuffer::FromCompressedNoValidate(std::move(ChunkBuffer));
				ZEN_ASSERT(Chunk);
				Attachments.push_back(zen::CbAttachment(Chunk, ChunkEntry.ChunkHash));
			}
			else
			{
				uint32_t BlockIndex = gsl::narrow<uint32_t>(ChunkEntry.Offset / BlockSize);
				if (BlockIndex != ReadBlockIndex)
				{
					ReadBlockIndex = BlockIndex;
					BlockStore	   = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksPath(m_SourcePath, ReadBlockIndex));
					ZEN_ASSERT(BlockStore);
				}
				size_t				  BlockOffset		   = BlockIndex * BlockSize;
				size_t				  AttachmentBulkOffset = ChunkEntry.Offset - BlockOffset;
				zen::IoBuffer		  ChunkBuffer(BlockStore, AttachmentBulkOffset, ChunkEntry.Length);
				zen::CompressedBuffer Chunk = zen::CompressedBuffer::FromCompressedNoValidate(std::move(ChunkBuffer));
				Attachments.push_back(zen::CbAttachment(Chunk, ChunkEntry.ChunkHash));
			}
		}
		RequestPackage.AddAttachments(Attachments);

		ZEN_INFO("Sending oplog with {} ops and {} attachments for '{}/{}' to {}",
				 OpEntries.size(),
				 ChunksHeader->ChunkCount,
				 m_ProjectName,
				 OplogName,
				 m_HostName);
		std::vector<zen::IoBuffer>	   RequestPayload = zen::FormatPackageMessage(RequestPackage, zen::FormatFlags::kAllowLocalReferences);
		std::vector<zen::SharedBuffer> Parts;
		Parts.reserve(RequestPayload.size());
		for (const auto& I : RequestPayload)
		{
			Parts.emplace_back(zen::SharedBuffer(I));
		}
		zen::CompositeBuffer  Cmp(std::move(Parts));
		zen::CompressedBuffer CompressedRequest = zen::CompressedBuffer::Compress(Cmp);

		std::string AppendOplogRequest = fmt::format("{}/{}/oplog/{}/archive", UrlBase, m_ProjectName, OplogName);
		Session.SetUrl(AppendOplogRequest);

		zen::IoBuffer TmpBuffer = CompressedRequest.GetCompressed().Flatten().AsIoBuffer();
		Session.SetBody(cpr::Body{(const char*)TmpBuffer.GetData(), TmpBuffer.GetSize()});
		cpr::Response Response = Session.Post();
		if (!ExportProjectCommand::IsSuccess(Response, AppendOplogRequest))
		{
			return 1;
		}

		ZEN_INFO("Imported {} ops and {} chunks", OpEntries.size(), ChunksHeader->ChunkCount);
	}
	return 0;
}