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

#include "cachestore.h"

#include <zencore/crc32.h>
#include <zencore/except.h>
#include <zencore/logging.h>
#include <zencore/windows.h>

#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/string.h>
#include <zencore/thread.h>
#include <zenstore/basicfile.h>
#include <zenstore/cas.h>
#include <zenstore/caslog.h>

#include <fmt/core.h>
#include <concepts>
#include <filesystem>
#include <gsl/gsl-lite.hpp>
#include <unordered_map>

#include <atlfile.h>

using namespace zen;
using namespace fmt::literals;

namespace UE {

struct CorruptionTrailer
{
	enum
	{
		/** Arbitrary number used to identify corruption **/
		MagicConstant = 0x1e873d89
	};

	uint32_t Magic		   = MagicConstant;
	uint32_t Version	   = 1;
	uint32_t CRCofPayload  = 0;
	uint32_t SizeOfPayload = 0;

	void Initialize(const void* Data, size_t Size)
	{
		CRCofPayload  = zen::MemCrc32_Deprecated(Data, Size);
		SizeOfPayload = (uint32_t)Size;
	}
};

std::filesystem::path
GenerateDdcPath(std::string_view Key, std::filesystem::path& rootDir)
{
	std::filesystem::path FilePath = rootDir;

	std::string k8{Key};
	for (auto& c : k8)
		c = (char)toupper(c);

	const uint32_t Hash = zen::StrCrc_Deprecated(k8.c_str());

	std::wstring DirName;

	DirName = u'0' + ((Hash / 100) % 10);
	FilePath /= DirName;
	DirName = u'0' + ((Hash / 10) % 10);
	FilePath /= DirName;
	DirName = u'0' + (Hash % 10);
	FilePath /= DirName;

	FilePath /= Key;

	auto NativePath = FilePath.native();
	NativePath.append(L".udd");

	return NativePath;
}

}  // namespace UE

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

FileCacheStore::FileCacheStore(const char* RootDir, const char* ReadRootDir)
{
	// Ensure root directory exists - create if it doesn't exist already

	ZEN_INFO("Initializing FileCacheStore at '{}'", std::string_view(RootDir));

	m_RootDir = RootDir;

	std::error_code ErrorCode;

	std::filesystem::create_directories(m_RootDir, ErrorCode);

	if (ErrorCode)
	{
		ExtendableStringBuilder<256> Name;
		WideToUtf8(m_RootDir.c_str(), Name);

		ZEN_ERROR("Could not open file cache directory '{}' for writing ({})", Name.c_str(), ErrorCode.message());

		m_IsOk = false;
	}

	if (ReadRootDir)
	{
		m_ReadRootDir = ReadRootDir;

		if (std::filesystem::exists(m_ReadRootDir, ErrorCode))
		{
			ZEN_INFO("FileCacheStore will use additional read tree at '{}'", std::string_view(ReadRootDir));

			m_ReadRootIsValid = true;
		}
	}
}

FileCacheStore::~FileCacheStore()
{
}

bool
FileCacheStore::Get(std::string_view Key, CacheValue& OutValue)
{
	CAtlFile File;

	std::filesystem::path NativePath;

	HRESULT hRes = E_FAIL;

	if (m_ReadRootDir.empty() == false)
	{
		NativePath = UE::GenerateDdcPath(Key, m_ReadRootDir);

		hRes = File.Create(NativePath.c_str(), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING);
	}

	if (FAILED(hRes))
	{
		NativePath = UE::GenerateDdcPath(Key, m_RootDir);

		hRes = File.Create(NativePath.c_str(), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING);
	}

	if (FAILED(hRes))
	{
		ZEN_DEBUG("GET MISS {}", Key);

		return false;
	}

	ULONGLONG FileSize;
	File.GetSize(FileSize);

	if (FileSize <= 16)
	{
		return false;
	}

	FileSize -= 16;	 // CorruptionWrapper trailer

	OutValue.Value = IoBuffer(IoBuffer::File, File.Detach(), 0, FileSize);

	ZEN_DEBUG("GET HIT {}", Key);

	return true;
}

void
FileCacheStore::Put(std::string_view Key, const CacheValue& Value)
{
	const void* Data = Value.Value.Data();
	size_t		Size = Value.Value.Size();

	UE::CorruptionTrailer Trailer;
	Trailer.Initialize(Data, Size);

	std::filesystem::path NativePath = UE::GenerateDdcPath(Key, m_RootDir);

	CAtlTemporaryFile File;

	ZEN_DEBUG("PUT {}", Key);

	HRESULT hRes = File.Create(m_RootDir.c_str());

	if (SUCCEEDED(hRes))
	{
		const uint8_t* WritePointer = reinterpret_cast<const uint8_t*>(Data);

		while (Size)
		{
			const int MaxChunkSize = 16 * 1024 * 1024;
			const int ChunkSize	   = (int)((Size > MaxChunkSize) ? MaxChunkSize : Size);

			DWORD BytesWritten = 0;
			File.Write(WritePointer, ChunkSize, &BytesWritten);

			Size -= BytesWritten;
			WritePointer += BytesWritten;
		}

		File.Write(&Trailer, sizeof Trailer);
		hRes = File.Close(NativePath.c_str());	// This renames the file to its final name

		if (FAILED(hRes))
		{
			ZEN_WARN("Failed to rename temp file for key '{}' - deleting temporary file", Key);

			if (!DeleteFile(File.TempFileName()))
			{
				ZEN_WARN("Temp file for key '{}' could not be deleted - no value persisted", Key);
			}
		}
	}
}

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

MemoryCacheStore::MemoryCacheStore()
{
}

MemoryCacheStore::~MemoryCacheStore()
{
}

bool
MemoryCacheStore::Get(std::string_view InKey, CacheValue& OutValue)
{
	RwLock::SharedLockScope _(m_Lock);

	auto it = m_CacheMap.find(std::string(InKey));

	if (it == m_CacheMap.end())
	{
		return false;
	}
	else
	{
		OutValue.Value = it->second;

		return true;
	}
}

void
MemoryCacheStore::Put(std::string_view Key, const CacheValue& Value)
{
	RwLock::ExclusiveLockScope _(m_Lock);
	m_CacheMap[std::string(Key)] = Value.Value;
}