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

#include <objectstore/objectstore.h>

#include <zencore/base64.h>
#include <zencore/compactbinaryvalue.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/string.h>
#include "zencore/compactbinary.h"
#include "zencore/compactbinarybuilder.h"
#include "zenhttp/httpcommon.h"
#include "zenhttp/httpserver.h"

#include <filesystem>
#include <thread>

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

using namespace std::literals;

ZEN_DEFINE_LOG_CATEGORY_STATIC(LogObj, "obj"sv);

class CbXmlWriter
{
public:
	explicit CbXmlWriter(StringBuilderBase& InBuilder) : Builder(InBuilder)
	{
		Builder.Append("<?xml version=\"1.0\" encoding=\"UTF-8\"?>");
		Builder << LINE_TERMINATOR_ANSI;
	}

	void WriteField(CbFieldView Field)
	{
		using namespace std::literals;

		bool					 SkipEndTag = false;
		const std::u8string_view Tag		= Field.GetU8Name();

		AppendBeginTag(Tag);

		switch (CbValue Accessor = Field.GetValue(); Accessor.GetType())
		{
			case CbFieldType::Null:
				Builder << "Null"sv;
				break;
			case CbFieldType::Object:
			case CbFieldType::UniformObject:
				{
					for (CbFieldView It : Field)
					{
						WriteField(It);
					}
				}
				break;
			case CbFieldType::Array:
			case CbFieldType::UniformArray:
				{
					bool FirstField = true;
					for (CbFieldView It : Field)
					{
						if (!FirstField)
							AppendBeginTag(Tag);

						WriteField(It);
						AppendEndTag(Tag);
						FirstField = false;
					}
					SkipEndTag = true;
				}
				break;
			case CbFieldType::Binary:
				AppendBase64String(Accessor.AsBinary());
				break;
			case CbFieldType::String:
				Builder << Accessor.AsU8String();
				break;
			case CbFieldType::IntegerPositive:
				Builder << Accessor.AsIntegerPositive();
				break;
			case CbFieldType::IntegerNegative:
				Builder << Accessor.AsIntegerNegative();
				break;
			case CbFieldType::Float32:
				{
					const float Value = Accessor.AsFloat32();
					if (std::isfinite(Value))
					{
						Builder.Append(fmt::format("{:.9g}", Value));
					}
					else
					{
						Builder << "Null"sv;
					}
				}
				break;
			case CbFieldType::Float64:
				{
					const double Value = Accessor.AsFloat64();
					if (std::isfinite(Value))
					{
						Builder.Append(fmt::format("{:.17g}", Value));
					}
					else
					{
						Builder << "null"sv;
					}
				}
				break;
			case CbFieldType::BoolFalse:
				Builder << "False"sv;
				break;
			case CbFieldType::BoolTrue:
				Builder << "True"sv;
				break;
			case CbFieldType::ObjectAttachment:
			case CbFieldType::BinaryAttachment:
				{
					Accessor.AsAttachment().ToHexString(Builder);
				}
				break;
			case CbFieldType::Hash:
				{
					Accessor.AsHash().ToHexString(Builder);
				}
				break;
			case CbFieldType::Uuid:
				{
					Accessor.AsUuid().ToString(Builder);
				}
				break;
			case CbFieldType::DateTime:
				Builder << DateTime(Accessor.AsDateTimeTicks()).ToIso8601();
				break;
			case CbFieldType::TimeSpan:
				{
					const TimeSpan Span(Accessor.AsTimeSpanTicks());
					if (Span.GetDays() == 0)
					{
						Builder << Span.ToString("%h:%m:%s.%n");
					}
					else
					{
						Builder << Span.ToString("%d.%h:%m:%s.%n");
					}
					break;
				}
			case CbFieldType::ObjectId:
				Accessor.AsObjectId().ToString(Builder);
				break;
			case CbFieldType::CustomById:
				{
					CbCustomById Custom = Accessor.AsCustomById();

					AppendBeginTag(u8"Id"sv);
					Builder << Custom.Id;
					AppendEndTag(u8"Id"sv);

					AppendBeginTag(u8"Data"sv);
					AppendBase64String(Custom.Data);
					AppendEndTag(u8"Data"sv);
					break;
				}
			case CbFieldType::CustomByName:
				{
					CbCustomByName Custom = Accessor.AsCustomByName();

					AppendBeginTag(u8"Name"sv);
					Builder << Custom.Name;
					AppendEndTag(u8"Name"sv);

					AppendBeginTag(u8"Data"sv);
					AppendBase64String(Custom.Data);
					AppendEndTag(u8"Data"sv);
					break;
				}
			default:
				ZEN_ASSERT(false);
				break;
		}

		if (!SkipEndTag)
			AppendEndTag(Tag);
	}

private:
	void AppendBeginTag(std::u8string_view Tag)
	{
		if (!Tag.empty())
		{
			Builder << '<' << Tag << '>';
		}
	}

	void AppendEndTag(std::u8string_view Tag)
	{
		if (!Tag.empty())
		{
			Builder << "</"sv << Tag << '>';
		}
	}

	void AppendBase64String(MemoryView Value)
	{
		Builder << '"';
		ZEN_ASSERT(Value.GetSize() <= 512 * 1024 * 1024);
		const uint32_t EncodedSize	= Base64::GetEncodedDataSize(uint32_t(Value.GetSize()));
		const size_t   EncodedIndex = Builder.AddUninitialized(size_t(EncodedSize));
		Base64::Encode(static_cast<const uint8_t*>(Value.GetData()), uint32_t(Value.GetSize()), Builder.Data() + EncodedIndex);
	}

private:
	StringBuilderBase& Builder;
};

HttpObjectStoreService::HttpObjectStoreService(ObjectStoreConfig Cfg) : m_Cfg(std::move(Cfg))
{
	Inititalize();
}

HttpObjectStoreService::~HttpObjectStoreService()
{
}

const char*
HttpObjectStoreService::BaseUri() const
{
	return "/obj/";
}

void
HttpObjectStoreService::HandleRequest(zen::HttpServerRequest& Request)
{
	if (m_Router.HandleRequest(Request) == false)
	{
		ZEN_LOG_WARN(LogObj, "No route found for {0}", Request.RelativeUri());
		return Request.WriteResponse(HttpResponseCode::NotFound, HttpContentType::kText, "Not found"sv);
	}
}

void
HttpObjectStoreService::Inititalize()
{
	namespace fs = std::filesystem;
	ZEN_LOG_INFO(LogObj, "Initialzing Object Store in '{}'", m_Cfg.RootDirectory);

	const fs::path BucketsPath = m_Cfg.RootDirectory / "buckets";
	if (!fs::exists(BucketsPath))
	{
		CreateDirectories(BucketsPath);
	}

	m_Router.RegisterRoute(
		"bucket",
		[this](zen::HttpRouterRequest& Request) { CreateBucket(Request); },
		HttpVerb::kPost | HttpVerb::kPut);

	m_Router.RegisterRoute(
		"bucket",
		[this](zen::HttpRouterRequest& Request) { DeleteBucket(Request); },
		HttpVerb::kDelete);

	m_Router.RegisterRoute(
		"bucket/{path}",
		[this](zen::HttpRouterRequest& Request) {
			const std::string Path	   = Request.GetCapture(1);
			const auto		  Sep	   = Path.find_last_of('.');
			const bool		  IsObject = Sep != std::string::npos && Path.size() - Sep > 0;

			if (IsObject)
			{
				GetObject(Request, Path);
			}
			else
			{
				ListBucket(Request, Path);
			}
		},
		HttpVerb::kHead | HttpVerb::kGet);

	m_Router.RegisterRoute(
		"bucket/{bucket}/{path}",
		[this](zen::HttpRouterRequest& Request) { PutObject(Request); },
		HttpVerb::kPost | HttpVerb::kPut);
}

std::filesystem::path
HttpObjectStoreService::GetBucketDirectory(std::string_view BucketName)
{
	{
		std::lock_guard _(BucketsMutex);

		if (const auto It = std::find_if(std::begin(m_Cfg.Buckets),
										 std::end(m_Cfg.Buckets),
										 [&BucketName](const auto& Bucket) -> bool { return Bucket.Name == BucketName; });
			It != std::end(m_Cfg.Buckets))
		{
			return It->Directory.make_preferred();
		}
	}

	return (m_Cfg.RootDirectory / "buckets" / BucketName).make_preferred();
}

void
HttpObjectStoreService::CreateBucket(zen::HttpRouterRequest& Request)
{
	namespace fs = std::filesystem;

	const CbObject		   Params	  = Request.ServerRequest().ReadPayloadObject();
	const std::string_view BucketName = Params["bucketname"].AsString();

	if (BucketName.empty())
	{
		return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
	}

	const fs::path BucketPath = m_Cfg.RootDirectory / "buckets" / BucketName;
	{
		std::lock_guard _(BucketsMutex);
		if (!fs::exists(BucketPath))
		{
			CreateDirectories(BucketPath);
			ZEN_LOG_INFO(LogObj, "CREATE - new bucket '{}' OK", BucketName);
			return Request.ServerRequest().WriteResponse(HttpResponseCode::Created);
		}
	}

	ZEN_LOG_INFO(LogObj, "CREATE - existing bucket '{}' OK", BucketName);
	Request.ServerRequest().WriteResponse(HttpResponseCode::OK);
}

void
HttpObjectStoreService::ListBucket(zen::HttpRouterRequest& Request, const std::string& Path)
{
	namespace fs = std::filesystem;

	const auto		  Sep		 = Path.find_first_of('/');
	const std::string BucketName = Sep == std::string::npos ? Path : Path.substr(0, Sep);
	if (BucketName.empty())
	{
		return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
	}

	std::string BucketPrefix = Sep == std::string::npos || Sep == Path.size() - 1 ? std::string() : Path.substr(BucketName.size() + 1);
	if (BucketPrefix.empty())
	{
		const auto QueryParms = Request.ServerRequest().GetQueryParams();
		if (auto PrefixParam = QueryParms.GetValue("prefix"); PrefixParam.empty() == false)
		{
			BucketPrefix = PrefixParam;
		}
	}
	BucketPrefix.erase(0, BucketPrefix.find_first_not_of('/'));
	BucketPrefix.erase(0, BucketPrefix.find_first_not_of('\\'));

	const fs::path BucketRoot		  = GetBucketDirectory(BucketName);
	const fs::path RelativeBucketPath = fs::path(BucketPrefix).make_preferred();
	const fs::path FullPath			  = BucketRoot / RelativeBucketPath;

	struct Visitor : FileSystemTraversal::TreeVisitor
	{
		Visitor(const std::string_view BucketName, const fs::path& Path, const fs::path& Prefix) : BucketPath(Path)
		{
			Writer.BeginObject("ListBucketResult"sv);
			Writer << "Name"sv << BucketName;
			std::string Tmp = Prefix.string();
			std::replace(Tmp.begin(), Tmp.end(), '\\', '/');
			Writer << "Prefix"sv << Tmp;
			Writer.BeginArray("Contents"sv);
		}

		void VisitFile(const fs::path& Parent, const path_view& File, uint64_t FileSize) override
		{
			const fs::path FullPath		= Parent / fs::path(File);
			fs::path	   RelativePath = fs::relative(FullPath, BucketPath);

			std::string Key = RelativePath.string();
			std::replace(Key.begin(), Key.end(), '\\', '/');

			Writer.BeginObject();
			Writer << "Key"sv << Key;
			Writer << "Size"sv << FileSize;
			Writer.EndObject();
		}

		bool VisitDirectory(const std::filesystem::path&, const path_view&) override { return false; }

		CbObject GetResult()
		{
			Writer.EndArray();
			Writer.EndObject();
			return Writer.Save();
		}

		CbObjectWriter Writer;
		fs::path	   BucketPath;
	};

	Visitor				FileVisitor(BucketName, BucketRoot, RelativeBucketPath);
	FileSystemTraversal Traversal;

	if (std::filesystem::exists(FullPath))
	{
		std::lock_guard _(BucketsMutex);
		Traversal.TraverseFileSystem(FullPath, FileVisitor);
	}
	CbObject Result = FileVisitor.GetResult();

	if (Request.ServerRequest().AcceptContentType() == HttpContentType::kJSON)
	{
		ExtendableStringBuilder<1024> Sb;
		return Request.ServerRequest().WriteResponse(HttpResponseCode::OK, HttpContentType::kJSON, Result.ToJson(Sb).ToView());
	}

	ExtendableStringBuilder<1024> Xml;
	CbXmlWriter					  XmlWriter(Xml);
	XmlWriter.WriteField(Result.AsFieldView());

	Request.ServerRequest().WriteResponse(HttpResponseCode::OK, HttpContentType::kXML, Xml.ToView());
}

void
HttpObjectStoreService::DeleteBucket(zen::HttpRouterRequest& Request)
{
	namespace fs = std::filesystem;

	const CbObject		   Params	  = Request.ServerRequest().ReadPayloadObject();
	const std::string_view BucketName = Params["bucketname"].AsString();

	if (BucketName.empty())
	{
		return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
	}

	const fs::path BucketPath = m_Cfg.RootDirectory / "buckets" / BucketName;
	{
		std::lock_guard _(BucketsMutex);
		DeleteDirectories(BucketPath);
	}

	ZEN_LOG_INFO(LogObj, "DELETE - bucket '{}' OK", BucketName);
	Request.ServerRequest().WriteResponse(HttpResponseCode::OK);
}

void
HttpObjectStoreService::GetObject(zen::HttpRouterRequest& Request, const std::string& Path)
{
	namespace fs = std::filesystem;

	const auto		  Sep		 = Path.find_first_of('/');
	const std::string BucketName = Sep == std::string::npos ? Path : Path.substr(0, Sep);
	const std::string BucketPrefix =
		Sep == std::string::npos || Sep == Path.size() - 1 ? std::string() : Path.substr(BucketName.size() + 1);

	const fs::path BucketDir = GetBucketDirectory(BucketName);

	if (BucketDir.empty())
	{
		ZEN_LOG_DEBUG(LogObj, "GET - [FAILED], unknown bucket '{}'", BucketName);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::NotFound);
	}

	const fs::path RelativeBucketPath = fs::path(BucketPrefix).make_preferred();

	if (RelativeBucketPath.is_absolute() || RelativeBucketPath.string().starts_with(".."))
	{
		ZEN_LOG_DEBUG(LogObj, "GET - from bucket '{}' [FAILED], invalid file path", BucketName);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::Forbidden);
	}

	const fs::path FilePath = BucketDir / RelativeBucketPath;
	if (!fs::exists(FilePath))
	{
		ZEN_LOG_DEBUG(LogObj, "GET - '{}/{}' [FAILED], doesn't exist", BucketName, FilePath);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::NotFound);
	}

	zen::HttpRanges Ranges;
	if (Request.ServerRequest().TryGetRanges(Ranges); Ranges.size() > 1)
	{
		// Only a single range is supported
		return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
	}

	FileContents File;
	{
		std::lock_guard _(BucketsMutex);
		File = ReadFile(FilePath);
	}

	if (File.ErrorCode)
	{
		ZEN_LOG_WARN(LogObj,
					 "GET - '{}/{}' [FAILED] ('{}': {})",
					 BucketName,
					 FilePath,
					 File.ErrorCode.category().name(),
					 File.ErrorCode.value());

		return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
	}

	const IoBuffer& FileBuf = File.Data[0];

	if (Ranges.empty())
	{
		const uint64_t TotalServed = TotalBytesServed.fetch_add(FileBuf.Size()) + FileBuf.Size();

		ZEN_LOG_DEBUG(LogObj,
					  "GET - '{}/{}' ({}) [OK] (Served: {})",
					  BucketName,
					  RelativeBucketPath,
					  NiceBytes(FileBuf.Size()),
					  NiceBytes(TotalServed));

		Request.ServerRequest().WriteResponse(HttpResponseCode::OK, HttpContentType::kBinary, FileBuf);
	}
	else
	{
		const auto	   Range	   = Ranges[0];
		const uint64_t RangeSize   = 1 + (Range.End - Range.Start);
		const uint64_t TotalServed = TotalBytesServed.fetch_add(RangeSize) + RangeSize;

		ZEN_LOG_DEBUG(LogObj,
					  "GET - '{}/{}' (Range: {}-{}) ({}/{}) [OK] (Served: {})",
					  BucketName,
					  RelativeBucketPath,
					  Range.Start,
					  Range.End,
					  NiceBytes(RangeSize),
					  NiceBytes(FileBuf.Size()),
					  NiceBytes(TotalServed));

		MemoryView RangeView = FileBuf.GetView().Mid(Range.Start, RangeSize);
		if (RangeView.GetSize() != RangeSize)
		{
			return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
		}

		IoBuffer RangeBuf = IoBuffer(IoBuffer::Wrap, RangeView.GetData(), RangeView.GetSize());
		Request.ServerRequest().WriteResponse(HttpResponseCode::PartialContent, HttpContentType::kBinary, RangeBuf);
	}
}

void
HttpObjectStoreService::PutObject(zen::HttpRouterRequest& Request)
{
	namespace fs = std::filesystem;

	const std::string& BucketName = Request.GetCapture(1);
	const fs::path	   BucketDir  = GetBucketDirectory(BucketName);

	if (BucketDir.empty())
	{
		ZEN_LOG_DEBUG(LogObj, "PUT - [FAILED], unknown bucket '{}'", BucketName);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::NotFound);
	}

	const fs::path RelativeBucketPath = fs::path(Request.GetCapture(2)).make_preferred();

	if (RelativeBucketPath.is_absolute() || RelativeBucketPath.string().starts_with(".."))
	{
		ZEN_LOG_DEBUG(LogObj, "PUT - bucket '{}' [FAILED], invalid file path", BucketName);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::Forbidden);
	}

	const fs::path FilePath		 = BucketDir / RelativeBucketPath;
	const fs::path FileDirectory = FilePath.parent_path();

	{
		std::lock_guard _(BucketsMutex);

		if (!fs::exists(FileDirectory))
		{
			CreateDirectories(FileDirectory);
		}

		const IoBuffer FileBuf = Request.ServerRequest().ReadPayload();

		if (FileBuf.Size() == 0)
		{
			ZEN_LOG_DEBUG(LogObj, "PUT - '{}' [FAILED], empty file", FilePath);
			return Request.ServerRequest().WriteResponse(HttpResponseCode::BadRequest);
		}

		WriteFile(FilePath, FileBuf);
		ZEN_LOG_DEBUG(LogObj,
					  "PUT - '{}' [OK] ({})",
					  (fs::path(BucketName) / RelativeBucketPath).make_preferred(),
					  NiceBytes(FileBuf.Size()));
	}

	Request.ServerRequest().WriteResponse(HttpResponseCode::OK);
}

}  // namespace zen