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

#include <zencore/compactbinary.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinaryvalidation.h>
#include <zencore/compress.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/stream.h>
#include <zencore/timer.h>
#include <zenhttp/httpserver.h>
#include <zenstore/CAS.h>

#include "cachekey.h"
#include "monitoring/httpstats.h"
#include "structuredcache.h"
#include "structuredcachestore.h"
#include "upstream/jupiter.h"
#include "upstream/upstreamcache.h"
#include "upstream/zen.h"
#include "zenstore/cidstore.h"

#include <zencore/compactbinarypackage.h>

#include <algorithm>
#include <atomic>
#include <filesystem>
#include <queue>
#include <thread>

#include <gsl/gsl-lite.hpp>

namespace zen {

using namespace std::literals;

//////////////////////////////////////////////////////////////////////////

CachePolicy
ParseCachePolicy(const HttpServerRequest::QueryParams& QueryParams)
{
	const CachePolicy QueryPolicy = zen::ParseQueryCachePolicy(QueryParams.GetValue("query"sv));
	const CachePolicy StorePolicy = zen::ParseStoreCachePolicy(QueryParams.GetValue("store"sv));
	const CachePolicy SkipPolicy  = zen::ParseSkipCachePolicy(QueryParams.GetValue("skip"sv));

	return QueryPolicy | StorePolicy | SkipPolicy;
}

//////////////////////////////////////////////////////////////////////////

HttpStructuredCacheService::HttpStructuredCacheService(ZenCacheStore&				  InCacheStore,
													   CasStore&					  InStore,
													   CidStore&					  InCidStore,
													   HttpStatsService&			  StatsService,
													   HttpStatusService&			  StatusService,
													   std::unique_ptr<UpstreamCache> UpstreamCache)
: m_Log(logging::Get("cache"))
, m_CacheStore(InCacheStore)
, m_StatsService(StatsService)
, m_StatusService(StatusService)
, m_CasStore(InStore)
, m_CidStore(InCidStore)
, m_UpstreamCache(std::move(UpstreamCache))
{
	m_StatsService.RegisterHandler("z$", *this);
	m_StatusService.RegisterHandler("z$", *this);
}

HttpStructuredCacheService::~HttpStructuredCacheService()
{
	ZEN_INFO("closing structured cache");

	m_StatsService.UnregisterHandler("z$", *this);
	m_StatusService.UnregisterHandler("z$", *this);
}

const char*
HttpStructuredCacheService::BaseUri() const
{
	return "/z$/";
}

void
HttpStructuredCacheService::Flush()
{
}

void
HttpStructuredCacheService::Scrub(ScrubContext& Ctx)
{
	if (m_LastScrubTime == Ctx.ScrubTimestamp())
	{
		return;
	}

	m_LastScrubTime = Ctx.ScrubTimestamp();

	m_CasStore.Scrub(Ctx);
	m_CidStore.Scrub(Ctx);
	m_CacheStore.Scrub(Ctx);
}

void
HttpStructuredCacheService::HandleRequest(HttpServerRequest& Request)
{
	CacheRef Ref;

	metrics::OperationTiming::Scope $(m_HttpRequests);

	if (!ValidateKeyUri(Request, /* out */ Ref))
	{
		std::string_view Key = Request.RelativeUri();

		if (Key == "$batch")
		{
			return HandleBatchRequest(Request);
		}

		if (std::all_of(begin(Key), end(Key), [](const char c) { return std::isalnum(c); }))
		{
			// Bucket reference

			return HandleCacheBucketRequest(Request, Key);
		}

		return Request.WriteResponse(HttpResponseCode::BadRequest);	 // invalid URL
	}

	const auto	QueryParams = Request.GetQueryParams();
	CachePolicy Policy		= ParseCachePolicy(QueryParams);

	if (Ref.PayloadId == IoHash::Zero)
	{
		return HandleCacheRecordRequest(Request, Ref, Policy);
	}
	else
	{
		return HandleCachePayloadRequest(Request, Ref, Policy);
	}

	return;
}

void
HttpStructuredCacheService::HandleCacheBucketRequest(HttpServerRequest& Request, std::string_view Bucket)
{
	ZEN_UNUSED(Request, Bucket);
	switch (auto Verb = Request.RequestVerb())
	{
		using enum HttpVerb;

		case kHead:
		case kGet:
			{
				// Query stats
			}
			break;

		case kDelete:
			// Drop bucket

			if (m_CacheStore.DropBucket(Bucket))
			{
				return Request.WriteResponse(HttpResponseCode::OK);
			}
			else
			{
				return Request.WriteResponse(HttpResponseCode::NotFound);
			}
			break;
	}
}

void
HttpStructuredCacheService::HandleCacheRecordRequest(HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	switch (auto Verb = Request.RequestVerb())
	{
		using enum HttpVerb;

		case kHead:
		case kGet:
			{
				HandleGetCacheRecord(Request, Ref, Policy);
			}
			break;

		case kPut:
			HandlePutCacheRecord(Request, Ref, Policy);
			break;
		default:
			break;
	}
}

void
HttpStructuredCacheService::HandleGetCacheRecord(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	const ZenContentType AcceptType		 = Request.AcceptContentType();
	const bool			 SkipData		 = (Policy & CachePolicy::SkipData) == CachePolicy::SkipData;
	const bool			 SkipAttachments = (Policy & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
	const bool			 QueryUpstream	 = m_UpstreamCache && ((Policy & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);

	bool		  Success = false;
	ZenCacheValue LocalCacheValue;

	if (m_CacheStore.Get(Ref.BucketSegment, Ref.HashKey, LocalCacheValue))
	{
		Success = true;

		if (!SkipData && AcceptType == ZenContentType::kCbPackage)
		{
			CbPackage	 Package;
			CbObjectView CacheRecord(LocalCacheValue.Value.Data());
			uint32_t	 AttachmentCount = 0;
			uint32_t	 ValidCount		 = 0;

			if (!SkipAttachments)
			{
				CacheRecord.IterateAttachments([this, &Ref, &Package, &AttachmentCount, &ValidCount](CbFieldView AttachmentHash) {
					if (IoBuffer Chunk = m_CidStore.FindChunkByCid(AttachmentHash.AsHash()))
					{
						Package.AddAttachment(CbAttachment(CompressedBuffer::FromCompressed(SharedBuffer(Chunk))));
						ValidCount++;
					}
					AttachmentCount++;
				});

				if (ValidCount != AttachmentCount)
				{
					// Success = false;
					ZEN_WARN("GET - '{}/{}' '{}' is partial, found '{}' of '{}' attachments",
							 Ref.BucketSegment,
							 Ref.HashKey,
							 ToString(AcceptType),
							 ValidCount,
							 AttachmentCount);
				}
			}

			Package.SetObject(LoadCompactBinaryObject(LocalCacheValue.Value));

			BinaryWriter MemStream;
			Package.Save(MemStream);

			LocalCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
			LocalCacheValue.Value.SetContentType(HttpContentType::kCbPackage);
		}
	}

	if (Success)
	{
		ZEN_DEBUG("HIT - '{}/{}' {} '{}' (LOCAL)",
				  Ref.BucketSegment,
				  Ref.HashKey,
				  NiceBytes(LocalCacheValue.Value.Size()),
				  ToString(LocalCacheValue.Value.GetContentType()));

		m_CacheStats.HitCount++;

		if (SkipData)
		{
			return Request.WriteResponse(HttpResponseCode::OK);
		}
		else
		{
			return Request.WriteResponse(HttpResponseCode::OK, LocalCacheValue.Value.GetContentType(), LocalCacheValue.Value);
		}
	}
	else if (!QueryUpstream)
	{
		ZEN_DEBUG("MISS - '{}/{}' '{}'", Ref.BucketSegment, Ref.HashKey, ToString(AcceptType));
		m_CacheStats.MissCount++;
		return Request.WriteResponse(HttpResponseCode::NotFound);
	}

	// Issue upstream query asynchronously in order to keep requests flowing without
	// hogging I/O servicing threads with blocking work

	Request.WriteResponseAsync([this, AcceptType, SkipData, SkipAttachments, Ref](HttpServerRequest& AsyncRequest) {
		bool		  Success = false;
		ZenCacheValue UpstreamCacheValue;

		metrics::OperationTiming::Scope $(m_UpstreamGetRequestTiming);

		if (GetUpstreamCacheResult UpstreamResult = m_UpstreamCache->GetCacheRecord({Ref.BucketSegment, Ref.HashKey}, AcceptType);
			UpstreamResult.Success)
		{
			Success					 = true;
			UpstreamCacheValue.Value = UpstreamResult.Value;

			UpstreamCacheValue.Value.SetContentType(AcceptType);

			if (AcceptType == ZenContentType::kBinary || AcceptType == ZenContentType::kCbObject)
			{
				if (AcceptType == ZenContentType::kCbObject)
				{
					const CbValidateError ValidationResult = ValidateCompactBinary(UpstreamResult.Value, CbValidateMode::All);

					if (ValidationResult == CbValidateError::None)
					{
						CbObjectView CacheRecord(UpstreamResult.Value.Data());

						CbObjectWriter IndexData;
						IndexData.BeginArray("references");
						CacheRecord.IterateAttachments([&](CbFieldView Attachment) { IndexData.AddHash(Attachment.AsHash()); });
						IndexData.EndArray();

						UpstreamCacheValue.IndexData = IndexData.Save();
					}
					else
					{
						Success = false;
						ZEN_WARN("Get - '{}/{}' '{}' FAILED, invalid compact binary object from upstream",
								 Ref.BucketSegment,
								 Ref.HashKey,
								 ToString(AcceptType));
					}
				}

				if (Success)
				{
					m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, UpstreamCacheValue);
				}
			}
			else if (AcceptType == ZenContentType::kCbPackage)
			{
				CbPackage Package;
				if (Package.TryLoad(UpstreamCacheValue.Value))
				{
					CbObject CacheRecord	 = Package.GetObject();
					uint32_t AttachmentCount = 0;
					uint32_t ValidCount		 = 0;

					CacheRecord.IterateAttachments([this, &Package, &Ref, &AttachmentCount, &ValidCount](CbFieldView AttachmentHash) {
						if (const CbAttachment* Attachment = Package.FindAttachment(AttachmentHash.AsHash()))
						{
							if (CompressedBuffer Chunk = Attachment->AsCompressedBinary())
							{
								m_CidStore.AddChunk(Chunk);
								ValidCount++;
							}
							else
							{
								ZEN_WARN("Get - '{}/{}' '{}' FAILED, upstream attachment not compressed",
										 Ref.BucketSegment,
										 Ref.HashKey,
										 ToString(ZenContentType::kCbPackage));
							}
						}
						AttachmentCount++;
					});

					if (ValidCount == AttachmentCount)
					{
						m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = CacheRecord.GetBuffer().AsIoBuffer()});

						if (SkipAttachments)
						{
							CbPackage PackageWithoutAttachments;
							PackageWithoutAttachments.SetObject(CacheRecord);

							BinaryWriter MemStream;
							PackageWithoutAttachments.Save(MemStream);

							UpstreamCacheValue.Value = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
						}
					}
					else
					{
						Success = false;
						ZEN_WARN("Get - '{}/{}' '{}' FAILED, attachments missing in upstream package",
								 Ref.BucketSegment,
								 Ref.HashKey,
								 ToString(AcceptType));
					}
				}
				else
				{
					Success = false;
					ZEN_WARN("Get - '{}/{}' '{}' FAILED, invalid upstream package", Ref.BucketSegment, Ref.HashKey, ToString(AcceptType));
				}
			}
		}

		if (Success)
		{
			ZEN_DEBUG("HIT - '{}/{}' {} '{}' (UPSTREAM)",
					  Ref.BucketSegment,
					  Ref.HashKey,
					  NiceBytes(UpstreamCacheValue.Value.Size()),
					  ToString(UpstreamCacheValue.Value.GetContentType()));

			m_CacheStats.HitCount++;
			m_CacheStats.UpstreamHitCount++;

			if (SkipData)
			{
				AsyncRequest.WriteResponse(HttpResponseCode::OK);
			}
			else
			{
				AsyncRequest.WriteResponse(HttpResponseCode::OK, UpstreamCacheValue.Value.GetContentType(), UpstreamCacheValue.Value);
			}
		}
		else
		{
			ZEN_DEBUG("MISS - '{}/{}' '{}'", Ref.BucketSegment, Ref.HashKey, ToString(AcceptType));
			m_CacheStats.MissCount++;
			AsyncRequest.WriteResponse(HttpResponseCode::NotFound);
		}
	});
}

void
HttpStructuredCacheService::HandlePutCacheRecord(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	IoBuffer Body = Request.ReadPayload();

	if (!Body || Body.Size() == 0)
	{
		return Request.WriteResponse(HttpResponseCode::BadRequest);
	}

	const HttpContentType ContentType	= Request.RequestContentType();
	const bool			  StoreUpstream = m_UpstreamCache && (CachePolicy::StoreRemote == (Policy & CachePolicy::StoreRemote));

	Body.SetContentType(ContentType);

	if (ContentType == HttpContentType::kBinary)
	{
		ZEN_DEBUG("PUT - '{}/{}' {} '{}'", Ref.BucketSegment, Ref.HashKey, NiceBytes(Body.Size()), ToString(ContentType));
		m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = Body});

		if (StoreUpstream)
		{
			ZEN_ASSERT(m_UpstreamCache);
			auto Result = m_UpstreamCache->EnqueueUpstream({.Type = ZenContentType::kBinary, .CacheKey = {Ref.BucketSegment, Ref.HashKey}});
		}

		Request.WriteResponse(HttpResponseCode::Created);
	}
	else if (ContentType == HttpContentType::kCbObject)
	{
		const CbValidateError ValidationResult = ValidateCompactBinary(MemoryView(Body.GetData(), Body.GetSize()), CbValidateMode::All);

		if (ValidationResult != CbValidateError::None)
		{
			ZEN_WARN("PUT - '{}/{}' '{}' FAILED, invalid compact binary", Ref.BucketSegment, Ref.HashKey, ToString(ContentType));
			return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Compact binary validation failed"sv);
		}

		CbObjectView		CacheRecord(Body.Data());
		std::vector<IoHash> ValidAttachments;
		uint32_t			AttachmentCount = 0;

		CacheRecord.IterateAttachments([this, &AttachmentCount, &ValidAttachments](CbFieldView AttachmentHash) {
			const IoHash Hash = AttachmentHash.AsHash();
			if (m_CidStore.ContainsChunk(Hash))
			{
				ValidAttachments.emplace_back(Hash);
			}
			AttachmentCount++;
		});

		const uint32_t ValidCount		= static_cast<uint32_t>(ValidAttachments.size());
		const bool	   ValidCacheRecord = ValidCount == AttachmentCount;

		if (ValidCacheRecord)
		{
			ZEN_DEBUG("PUT - '{}/{}' {} '{}', {} attachments",
					  Ref.BucketSegment,
					  Ref.HashKey,
					  NiceBytes(Body.Size()),
					  ToString(ContentType),
					  ValidCount);

			m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = Body});

			if (StoreUpstream)
			{
				ZEN_ASSERT(m_UpstreamCache);
				auto Result = m_UpstreamCache->EnqueueUpstream({.Type		= ZenContentType::kCbObject,
																.CacheKey	= {Ref.BucketSegment, Ref.HashKey},
																.PayloadIds = std::move(ValidAttachments)});
			}

			Request.WriteResponse(HttpResponseCode::Created);
		}
		else
		{
			ZEN_WARN("PUT - '{}/{}' '{}' FAILED, found {}/{} attachments",
					 Ref.BucketSegment,
					 Ref.HashKey,
					 ToString(ContentType),
					 ValidCount,
					 AttachmentCount);

			Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Missing attachments"sv);
		}
	}
	else if (ContentType == HttpContentType::kCbPackage)
	{
		CbPackage Package;

		if (!Package.TryLoad(Body))
		{
			ZEN_WARN("PUT - '{}/{}' '{}' FAILED, invalid package", Ref.BucketSegment, Ref.HashKey, ToString(ContentType));
			return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Invalid package"sv);
		}

		CbObject CacheRecord = Package.GetObject();

		std::span<const CbAttachment> Attachments = Package.GetAttachments();
		std::vector<IoHash>			  ValidAttachments;
		int32_t						  NewAttachmentCount = 0;

		ValidAttachments.reserve(Attachments.size());

		CacheRecord.IterateAttachments([this, &Ref, &Package, &ValidAttachments, &NewAttachmentCount](CbFieldView AttachmentHash) {
			if (const CbAttachment* Attachment = Package.FindAttachment(AttachmentHash.AsHash()))
			{
				if (Attachment->IsCompressedBinary())
				{
					CompressedBuffer	   Chunk		= Attachment->AsCompressedBinary();
					CidStore::InsertResult InsertResult = m_CidStore.AddChunk(Chunk);

					ValidAttachments.emplace_back(InsertResult.DecompressedId);

					if (InsertResult.New)
					{
						NewAttachmentCount++;
					}
				}
				else
				{
					ZEN_WARN("PUT - '{}/{}' '{}' FAILED, attachment '{}' is not compressed",
							 Ref.BucketSegment,
							 Ref.HashKey,
							 ToString(HttpContentType::kCbPackage),
							 AttachmentHash.AsHash());
				}
			}
			else
			{
				ZEN_WARN("PUT - '{}/{}' '{}' FAILED, missing attachment '{}'",
						 Ref.BucketSegment,
						 Ref.HashKey,
						 ToString(HttpContentType::kCbPackage),
						 AttachmentHash.AsHash());
			}
		});

		const bool IsPartialCacheRecord = ValidAttachments.size() != Attachments.size();

		// if (!AttachmentsValid)
		//{
		//	return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Invalid attachments"sv);
		//}

		ZEN_DEBUG("PUT - '{}/{}' {} '{}', {}/{} new attachments",
				  Ref.BucketSegment,
				  Ref.HashKey,
				  NiceBytes(Body.GetSize()),
				  ToString(ContentType),
				  NewAttachmentCount,
				  Attachments.size());

		IoBuffer CacheRecordValue = CacheRecord.GetBuffer().AsIoBuffer();
		CacheRecordValue.SetContentType(ZenContentType::kCbObject);

		m_CacheStore.Put(Ref.BucketSegment, Ref.HashKey, {.Value = CacheRecord.GetBuffer().AsIoBuffer()});

		if (StoreUpstream && !IsPartialCacheRecord)
		{
			ZEN_ASSERT(m_UpstreamCache);
			auto Result = m_UpstreamCache->EnqueueUpstream({.Type		= ZenContentType::kCbPackage,
															.CacheKey	= {Ref.BucketSegment, Ref.HashKey},
															.PayloadIds = std::move(ValidAttachments)});
		}

		Request.WriteResponse(HttpResponseCode::Created);
	}
	else
	{
		return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Content-Type invalid"sv);
	}
}

void
HttpStructuredCacheService::HandleCachePayloadRequest(HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	switch (auto Verb = Request.RequestVerb())
	{
		using enum HttpVerb;

		case kHead:
		case kGet:
			{
				HandleGetCachePayload(Request, Ref, Policy);
			}
			break;
		case kPut:
			HandlePutCachePayload(Request, Ref, Policy);
			break;
		default:
			break;
	}
}

void
HttpStructuredCacheService::HandleGetCachePayload(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	IoBuffer   Payload		   = m_CidStore.FindChunkByCid(Ref.PayloadId);
	bool	   InUpstreamCache = false;
	const bool QueryUpstream   = !Payload && m_UpstreamCache && (CachePolicy::QueryRemote == (Policy & CachePolicy::QueryRemote));

	if (QueryUpstream)
	{
		if (auto UpstreamResult = m_UpstreamCache->GetCachePayload({Ref.BucketSegment, Ref.HashKey}, Ref.PayloadId);
			UpstreamResult.Success)
		{
			if (CompressedBuffer Compressed = CompressedBuffer::FromCompressed(SharedBuffer(UpstreamResult.Value)))
			{
				Payload							 = UpstreamResult.Value;
				IoHash				   ChunkHash = IoHash::HashBuffer(Payload);
				CasStore::InsertResult Result	 = m_CasStore.InsertChunk(Payload, ChunkHash);
				InUpstreamCache					 = true;

				m_CidStore.AddCompressedCid(Ref.PayloadId, ChunkHash);
			}
			else
			{
				ZEN_WARN("got uncompressed upstream cache payload");
			}
		}
	}

	if (!Payload)
	{
		ZEN_DEBUG("MISS - '{}/{}/{}' '{}'", Ref.BucketSegment, Ref.HashKey, Ref.PayloadId, ToString(Request.AcceptContentType()));
		m_CacheStats.MissCount++;
		return Request.WriteResponse(HttpResponseCode::NotFound);
	}

	ZEN_DEBUG("HIT - '{}/{}/{}' {} '{}' ({})",
			  Ref.BucketSegment,
			  Ref.HashKey,
			  Ref.PayloadId,
			  NiceBytes(Payload.Size()),
			  ToString(Payload.GetContentType()),
			  InUpstreamCache ? "UPSTREAM" : "LOCAL");

	m_CacheStats.HitCount++;
	if (InUpstreamCache)
	{
		m_CacheStats.UpstreamHitCount++;
	}

	if ((Policy & CachePolicy::SkipData) == CachePolicy::SkipData)
	{
		Request.WriteResponse(HttpResponseCode::OK);
	}
	else
	{
		Request.WriteResponse(HttpResponseCode::OK, HttpContentType::kBinary, Payload);
	}
}

void
HttpStructuredCacheService::HandlePutCachePayload(zen::HttpServerRequest& Request, const CacheRef& Ref, CachePolicy Policy)
{
	// Note: Individual cache payloads are not propagated upstream until a valid cache record has been stored
	ZEN_UNUSED(Policy);

	IoBuffer Body = Request.ReadPayload();

	if (!Body || Body.Size() == 0)
	{
		return Request.WriteResponse(HttpResponseCode::BadRequest);
	}

	Body.SetContentType(Request.RequestContentType());

	IoHash			 ChunkHash	= IoHash::HashBuffer(Body);
	CompressedBuffer Compressed = CompressedBuffer::FromCompressed(SharedBuffer(Body));

	if (!Compressed)
	{
		return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Attachments must be compressed"sv);
	}

	if (IoHash::FromBLAKE3(Compressed.GetRawHash()) != Ref.PayloadId)
	{
		return Request.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, "Payload ID does not match attachment hash"sv);
	}

	CasStore::InsertResult Result = m_CasStore.InsertChunk(Body, ChunkHash);

	m_CidStore.AddCompressedCid(Ref.PayloadId, ChunkHash);

	ZEN_DEBUG("PUT - '{}/{}/{}' {} '{}' ({})",
			  Ref.BucketSegment,
			  Ref.HashKey,
			  Ref.PayloadId,
			  NiceBytes(Body.Size()),
			  ToString(Body.GetContentType()),
			  Result.New ? "NEW" : "OLD");

	const HttpResponseCode ResponseCode = Result.New ? HttpResponseCode::Created : HttpResponseCode::OK;

	Request.WriteResponse(ResponseCode);
}

bool
HttpStructuredCacheService::ValidateKeyUri(HttpServerRequest& Request, CacheRef& OutRef)
{
	std::string_view			Key				  = Request.RelativeUri();
	std::string_view::size_type BucketSplitOffset = Key.find_first_of('/');

	if (BucketSplitOffset == std::string_view::npos)
	{
		return false;
	}

	OutRef.BucketSegment = ToLower(Key.substr(0, BucketSplitOffset));

	if (!std::all_of(begin(OutRef.BucketSegment), end(OutRef.BucketSegment), [](const char c) { return std::isalnum(c); }))
	{
		return false;
	}

	std::string_view HashSegment;
	std::string_view PayloadSegment;

	std::string_view::size_type PayloadSplitOffset = Key.find_last_of('/');

	// We know there is a slash so no need to check for npos return

	if (PayloadSplitOffset == BucketSplitOffset)
	{
		// Basic cache record lookup
		HashSegment = Key.substr(BucketSplitOffset + 1);
	}
	else
	{
		// Cache record + payload lookup
		HashSegment	   = Key.substr(BucketSplitOffset + 1, PayloadSplitOffset - BucketSplitOffset - 1);
		PayloadSegment = Key.substr(PayloadSplitOffset + 1);
	}

	if (HashSegment.size() != IoHash::StringLength)
	{
		return false;
	}

	if (!PayloadSegment.empty() && PayloadSegment.size() == IoHash::StringLength)
	{
		const bool IsOk = ParseHexBytes(PayloadSegment.data(), PayloadSegment.size(), OutRef.PayloadId.Hash);

		if (!IsOk)
		{
			return false;
		}
	}
	else
	{
		OutRef.PayloadId = IoHash::Zero;
	}

	const bool IsOk = ParseHexBytes(HashSegment.data(), HashSegment.size(), OutRef.HashKey.Hash);

	if (!IsOk)
	{
		return false;
	}

	return true;
}

void
HttpStructuredCacheService::HandleBatchRequest(zen::HttpServerRequest& Request)
{
	switch (auto Verb = Request.RequestVerb())
	{
		using enum HttpVerb;

		case kPost:
			{
				const HttpContentType ContentType = Request.RequestContentType();
				const HttpContentType AcceptType  = Request.AcceptContentType();

				if (ContentType != HttpContentType::kCbObject || AcceptType != HttpContentType::kCbPackage)
				{
					return Request.WriteResponse(HttpResponseCode::BadRequest);
				}

				CbObject			   BatchRequest = zen::LoadCompactBinaryObject(Request.ReadPayload());
				const std::string_view Method		= BatchRequest["Method"sv].AsString();

				if (Method == "GetCacheRecords"sv)
				{
					HandleBatchGetCacheRecords(Request, BatchRequest);
				}
				else if (Method == "GetCachePayloads"sv)
				{
					HandleBatchGetCachePayloads(Request, BatchRequest);
				}
				else
				{
					Request.WriteResponse(HttpResponseCode::BadRequest);
				}
			}
			break;
		default:
			Request.WriteResponse(HttpResponseCode::BadRequest);
			break;
	}
}

void
HttpStructuredCacheService::HandleBatchGetCacheRecords(zen::HttpServerRequest& Request, CbObjectView BatchRequest)
{
	using namespace fmt::literals;

	const std::string_view Method = BatchRequest["Method"sv].AsString();
	ZEN_ASSERT(Method == "GetCacheRecords"sv);

	CbObjectView Params = BatchRequest["Params"sv].AsObjectView();

	CacheRecordPolicy Policy;
	CacheRecordPolicy::Load(Params["Policy"sv].AsObjectView(), Policy);

	const bool SkipAttachments = (Policy.GetRecordPolicy() & CachePolicy::SkipAttachments) == CachePolicy::SkipAttachments;
	const bool QueryRemote	   = m_UpstreamCache && ((Policy.GetRecordPolicy() & CachePolicy::QueryRemote) == CachePolicy::QueryRemote);

	std::vector<CacheKey> CacheKeys;
	std::vector<IoBuffer> CacheValues;
	std::vector<IoBuffer> Payloads;
	std::vector<size_t>	  UpstreamRequests;
	std::vector<size_t>	  Missing;

	for (CbFieldView QueryView : Params["CacheKeys"sv])
	{
		CbObjectView Query = QueryView.AsObjectView();
		CacheKeys.push_back(CacheKey::Create(Query["Bucket"sv].AsString(), Query["Hash"sv].AsHash()));
	}

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

	CacheValues.resize(CacheKeys.size());

	if (!SkipAttachments)
	{
		Payloads.reserve(CacheKeys.size());
	}

	for (size_t KeyIndex = 0; const CacheKey& Key : CacheKeys)
	{
		ZenCacheValue CacheValue;
		if (m_CacheStore.Get(Key.Bucket, Key.Hash, CacheValue))
		{
			CbObjectView CacheRecord(CacheValue.Value.Data());

			if (!SkipAttachments)
			{
				CacheRecord.IterateAttachments([this, &Payloads](CbFieldView AttachmentHash) {
					if (IoBuffer Chunk = m_CidStore.FindChunkByCid(AttachmentHash.AsHash()))
					{
						Payloads.push_back(Chunk);
					}
				});
			}

			ZEN_DEBUG("HIT - '{}/{}' {} '{}' (LOCAL)",
					  Key.Bucket,
					  Key.Hash,
					  NiceBytes(CacheValue.Value.Size()),
					  ToString(CacheValue.Value.GetContentType()));

			CacheValues[KeyIndex] = CacheValue.Value;
			m_CacheStats.HitCount++;
		}
		else if (QueryRemote)
		{
			UpstreamRequests.push_back(KeyIndex);
		}
		else
		{
			Missing.push_back(KeyIndex);
		}

		++KeyIndex;
	}

	if (!UpstreamRequests.empty() && m_UpstreamCache)
	{
		const auto OnCacheRecordGetComplete =
			[this, &CacheKeys, &CacheValues, &Payloads, &Missing, SkipAttachments](CacheRecordGetCompleteParams&& Params) {
				if (Params.Record)
				{
					Params.Record.IterateAttachments([&](CbFieldView AttachmentHash) {
						if (const CbAttachment* Attachment = Params.Package.FindAttachment(AttachmentHash.AsHash()))
						{
							if (CompressedBuffer Chunk = Attachment->AsCompressedBinary())
							{
								m_CidStore.AddChunk(Chunk);

								if (!SkipAttachments)
								{
									Payloads.push_back(Chunk.GetCompressed().Flatten().AsIoBuffer());
								}
							}
						}
					});

					ZEN_DEBUG("HIT - '{}/{}' {} '{}' (UPSTREAM)",
							  Params.CacheKey.Bucket,
							  Params.CacheKey.Hash,
							  NiceBytes(Params.Record.GetView().GetSize()),
							  ToString(HttpContentType::kCbObject));

					CacheValues[Params.KeyIndex] = IoBufferBuilder::MakeCloneFromMemory(Params.Record.GetView());
					m_CacheStats.UpstreamHitCount++;
				}
				else
				{
					Missing.push_back(Params.KeyIndex);
				}
			};

		m_UpstreamCache->GetCacheRecords(CacheKeys, UpstreamRequests, Policy, std::move(OnCacheRecordGetComplete));
	}

	for (size_t KeyIndex : Missing)
	{
		const CacheKey& Key = CacheKeys[KeyIndex];
		ZEN_DEBUG("MISS - '{}/{}'", Key.Bucket, Key.Hash);
		m_CacheStats.MissCount++;
	}

	CbObjectWriter BatchResponse;

	BatchResponse.BeginArray("Result"sv);
	for (const IoBuffer& Value : CacheValues)
	{
		if (Value)
		{
			CbObjectView Record(Value.Data());
			BatchResponse << Record;
		}
		else
		{
			BatchResponse.AddNull();
		}
	}
	BatchResponse.EndArray();

	CbPackage Package;
	Package.SetObject(BatchResponse.Save());

	for (const IoBuffer& Payload : Payloads)
	{
		Package.AddAttachment(CbAttachment(CompressedBuffer::FromCompressed(SharedBuffer(Payload))));
	}

	BinaryWriter MemStream;
	Package.Save(MemStream);

	Request.WriteResponse(HttpResponseCode::OK,
						  HttpContentType::kCbPackage,
						  IoBuffer(IoBuffer::Wrap, MemStream.GetData(), MemStream.GetSize()));
}

void
HttpStructuredCacheService::HandleBatchGetCachePayloads(zen::HttpServerRequest& Request, CbObjectView BatchRequest)
{
	using namespace fmt::literals;

	const auto GetChunkIdFromPayloadId = [](CbObjectView Record, const Oid& PayloadId) -> IoHash {
		if (CbObjectView ValueObject = Record["Value"].AsObjectView())
		{
			const Oid Id = ValueObject["Id"].AsObjectId();
			if (Id == PayloadId)
			{
				return ValueObject["RawHash"sv].AsHash();
			}
		}

		for (CbFieldView AttachmentView : Record["Attachments"sv])
		{
			CbObjectView AttachmentObject = AttachmentView.AsObjectView();
			const Oid	 Id				  = AttachmentObject["Id"].AsObjectId();

			if (Id == PayloadId)
			{
				return AttachmentObject["RawHash"sv].AsHash();
			}
		}

		return IoHash::Zero;
	};

	const std::string_view Method = BatchRequest["Method"sv].AsString();
	ZEN_ASSERT(Method == "GetCachePayloads"sv);

	CbObjectView Params = BatchRequest["Params"sv].AsObjectView();

	std::vector<CacheChunkRequest> ChunkRequests;
	std::vector<size_t>			   UpstreamRequests;
	std::vector<size_t>			   Missing;

	for (CbFieldView ChunkView : Params["ChunkRequests"sv])
	{
		CbObjectView   Chunk		  = ChunkView.AsObjectView();
		CbObjectView   CacheKeyObject = Chunk["Key"sv].AsObjectView();
		const CacheKey Key			  = CacheKey::Create(CacheKeyObject["Bucket"sv].AsString(), CacheKeyObject["Hash"sv].AsHash());
		const IoHash   ChunkId		  = IoHash::Zero;
		const Oid	   PayloadId	  = Chunk["PayloadId"sv].AsObjectId();
		const uint64_t RawOffset	  = Chunk["RawoffSet"sv].AsUInt64();
		const uint64_t RawSize		  = Chunk["RawSize"sv].AsUInt64();
		const uint32_t ChunkPolicy	  = Chunk["Policy"sv].AsUInt32();

		ChunkRequests.emplace_back(Key, ChunkId, PayloadId, RawOffset, RawSize, static_cast<CachePolicy>(ChunkPolicy));
	}

	std::stable_sort(ChunkRequests.begin(), ChunkRequests.end());

	CacheKey CurrentKey = CacheKey::Empty;
	IoBuffer CurrentRecordBuffer;

	for (CacheChunkRequest& ChunkRequest : ChunkRequests)
	{
		if (ChunkRequest.Key != CurrentKey)
		{
			CurrentKey = ChunkRequest.Key;

			ZenCacheValue CacheValue;
			if (m_CacheStore.Get(ChunkRequest.Key.Bucket, ChunkRequest.Key.Hash, CacheValue))
			{
				CurrentRecordBuffer = CacheValue.Value;
			}
		}

		if (CurrentRecordBuffer)
		{
			ChunkRequest.ChunkId = GetChunkIdFromPayloadId(CbObjectView(CurrentRecordBuffer.GetData()), ChunkRequest.PayloadId);
		}
	}

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

	std::vector<IoBuffer> Chunks;
	Chunks.resize(ChunkRequests.size());

	for (size_t RequestIndex = 0; const CacheChunkRequest& ChunkRequest : ChunkRequests)
	{
		const bool QueryLocal  = (ChunkRequest.Policy & CachePolicy::QueryLocal) == CachePolicy::QueryLocal;
		const bool QueryRemote = (ChunkRequest.Policy & CachePolicy::QueryRemote) == CachePolicy::QueryRemote;

		if (QueryLocal)
		{
			if (IoBuffer Chunk = m_CidStore.FindChunkByCid(ChunkRequest.ChunkId))
			{
				ZEN_DEBUG("HIT - '{}/{}/{}' {} '{}' ({})",
						  ChunkRequest.Key.Bucket,
						  ChunkRequest.Key.Hash,
						  ChunkRequest.ChunkId,
						  NiceBytes(Chunk.Size()),
						  ToString(Chunk.GetContentType()),
						  "LOCAL");

				Chunks[RequestIndex] = Chunk;
				m_CacheStats.HitCount++;
			}
			else if (QueryRemote)
			{
				UpstreamRequests.push_back(RequestIndex);
			}
			else
			{
				Missing.push_back(RequestIndex);
			}
		}
		else
		{
			ZEN_DEBUG("SKIP - '{}/{}/{}'", ChunkRequest.Key.Bucket, ChunkRequest.Key.Hash, ChunkRequest.ChunkId);
		}

		++RequestIndex;
	}

	if (!UpstreamRequests.empty() && m_UpstreamCache)
	{
		const auto OnCachePayloadGetComplete = [this, &ChunkRequests, &Chunks, &Missing](CachePayloadGetCompleteParams&& Params) {
			if (Params.Payload)
			{
				CompressedBuffer Compressed	  = CompressedBuffer::FromCompressed(SharedBuffer(Params.Payload));
				auto			 InsertResult = m_CidStore.AddChunk(Compressed);

				ZEN_DEBUG("HIT - '{}/{}/{}' {} '{}' ({})",
						  Params.Request.Key.Bucket,
						  Params.Request.Key.Hash,
						  Params.Request.ChunkId,
						  NiceBytes(Params.Payload.GetSize()),
						  ToString(Params.Payload.GetContentType()),
						  "UPSTREAM");

				Chunks[Params.RequestIndex] = std::move(Params.Payload);

				m_CacheStats.HitCount++;
				m_CacheStats.UpstreamHitCount++;
			}
			else
			{
				Missing.push_back(Params.RequestIndex);
			}
		};

		m_UpstreamCache->GetCachePayloads(ChunkRequests, UpstreamRequests, std::move(OnCachePayloadGetComplete));
	}

	for (size_t RequestIndex : Missing)
	{
		const CacheChunkRequest& ChunkRequest = ChunkRequests[RequestIndex];
		ZEN_DEBUG("MISS - '{}/{}/{}'", ChunkRequest.Key.Bucket, ChunkRequest.Key.Hash, ChunkRequest.ChunkId);
		m_CacheStats.MissCount++;
	}

	CbPackage	   Package;
	CbObjectWriter BatchResponse;

	BatchResponse.BeginArray("Result"sv);

	for (size_t ChunkIndex = 0; ChunkIndex < Chunks.size(); ++ChunkIndex)
	{
		if (Chunks[ChunkIndex])
		{
			BatchResponse << ChunkRequests[ChunkIndex].ChunkId;
			Package.AddAttachment(CbAttachment(CompressedBuffer::FromCompressed(SharedBuffer(std::move(Chunks[ChunkIndex])))));
		}
		else
		{
			BatchResponse << IoHash::Zero;
		}
	}
	BatchResponse.EndArray();

	Package.SetObject(BatchResponse.Save());

	BinaryWriter MemStream;
	Package.Save(MemStream);

	Request.WriteResponse(HttpResponseCode::OK,
						  HttpContentType::kCbPackage,
						  IoBuffer(IoBuffer::Wrap, MemStream.GetData(), MemStream.GetSize()));
}

void
HttpStructuredCacheService::HandleStatsRequest(zen::HttpServerRequest& Request)
{
	CbObjectWriter Cbo;

	EmitSnapshot("requests", m_HttpRequests, Cbo);
	EmitSnapshot("upstream_gets", m_UpstreamGetRequestTiming, Cbo);

	const uint64_t HitCount			= m_CacheStats.HitCount;
	const uint64_t UpstreamHitCount = m_CacheStats.UpstreamHitCount;
	const uint64_t MissCount		= m_CacheStats.MissCount;
	const uint64_t TotalCount		= HitCount + MissCount;

	Cbo.BeginObject("cache");
	Cbo << "hits" << HitCount << "misses" << MissCount;
	Cbo << "hit_ratio" << (TotalCount > 0 ? (double(HitCount) / double(TotalCount)) : 0.0);
	Cbo << "upstream_hits" << m_CacheStats.UpstreamHitCount;
	Cbo << "upstream_ratio" << (HitCount > 0 ? (double(UpstreamHitCount) / double(HitCount)) : 0.0);
	Cbo.EndObject();

	if (m_UpstreamCache)
	{
		Cbo.BeginObject("upstream");
		m_UpstreamCache->GetStatus(Cbo);
		Cbo.EndObject();
	}

	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

void
HttpStructuredCacheService::HandleStatusRequest(zen::HttpServerRequest& Request)
{
	CbObjectWriter Cbo;
	Cbo << "ok" << true;
	Request.WriteResponse(HttpResponseCode::OK, Cbo.Save());
}

}  // namespace zen