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

#include <objectstore/objectstore.h>

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

#include <thread>

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>
#include <json11.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

using namespace std::literals;

ZEN_DEFINE_LOG_CATEGORY_STATIC(LogObj, "obj"sv);

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()
{
	ZEN_LOG_INFO(LogObj, "Initialzing Object Store in '{}'", m_Cfg.RootDirectory);
	for (const auto& Bucket : m_Cfg.Buckets)
	{
		ZEN_LOG_INFO(LogObj, " - bucket '{}' -> '{}'", Bucket.Name, Bucket.Directory);
	}

	m_Router.RegisterRoute(
		"distributionpoints/{bucket}",
		[this](zen::HttpRouterRequest& Request) {
			const std::string BucketName = Request.GetCapture(1);

			ExtendableStringBuilder<1024> Json;
			{
				CbObjectWriter Writer;
				Writer.BeginArray("distributions");
				Writer << fmt::format("http://localhost:{}/obj/{}", m_Cfg.ServerPort, BucketName);
				Writer.EndArray();
				Writer.Save().ToJson(Json);
			}

			Request.ServerRequest().WriteResponse(HttpResponseCode::OK, HttpContentType::kJSON, Json.ToString());
		},
		HttpVerb::kGet);

	m_Router.RegisterRoute(
		"{bucket}/{path}",
		[this](zen::HttpRouterRequest& Request) { GetBlob(Request); },
		HttpVerb::kGet);

	m_Router.RegisterRoute(
		"{bucket}/{path}",
		[this](zen::HttpRouterRequest& Request) { PutBlob(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;
	}

	return std::filesystem::path();
}

void
HttpObjectStoreService::GetBlob(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, "GET - [FAILED], unknown bucket '{}'", BucketName);
		return Request.ServerRequest().WriteResponse(HttpResponseCode::NotFound);
	}

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

	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);
	}

	fs::path FilePath = BucketDir / RelativeBucketPath;
	if (fs::exists(FilePath) == false)
	{
		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 = 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   = 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::PutBlob(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 = Request.GetCapture(2);

	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);
	}

	fs::path	   FilePath = BucketDir / RelativeBucketPath;
	const IoBuffer FileBuf	= Request.ServerRequest().ReadPayload();

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

	WriteFile(FilePath, FileBuf);
	ZEN_LOG_DEBUG(LogObj, "PUT - '{}/{}' [OK] ({})", BucketName, RelativeBucketPath, NiceBytes(FileBuf.Size()));
	Request.ServerRequest().WriteResponse(HttpResponseCode::OK);
}

}  // namespace zen