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

#include "cachetracking.h"

#include <zencore/filesystem.h>
#include <zencore/logging.h>
#include <zencore/string.h>

#include <zencore/testing.h>
#include <zencore/testutils.h>

ZEN_THIRD_PARTY_INCLUDES_START
#pragma comment(lib, "Rpcrt4.lib")	// RocksDB made me do this
#include <fmt/format.h>
#include <rocksdb/db.h>
#include <tsl/robin_map.h>
#include <tsl/robin_set.h>
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

using namespace fmt::literals;

namespace rocksdb = ROCKSDB_NAMESPACE;

static constinit auto Epoch = std::chrono::time_point<std::chrono::system_clock>{};

uint32_t
GetCurrentTimeStamp()
{
	auto Duration = std::chrono::system_clock::now() - Epoch;
	auto Minutes  = std::chrono::duration_cast<std::chrono::minutes>(Duration).count();

	return Minutes;
}

struct CacheAccessSnapshot
{
public:
	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
	{
		BucketTracker* Tracker = MapBucketToIndex(std::string(BucketSegment));

		Tracker->Track(HashKey);
	}

private:
	struct BucketTracker
	{
		RwLock				   Lock;
		tsl::robin_set<IoHash> AccessedKeys;

		void Track(const IoHash& HashKey)
		{
			RwLock::ExclusiveLockScope _(Lock);
			AccessedKeys.insert(HashKey);
		}
	};

	BucketTracker* MapBucketToIndex(const std::string& BucketName)
	{
		RwLock::SharedLockScope _(m_Lock);
		if (auto It = m_BucketMapping.find(BucketName); It == m_BucketMapping.end())
		{
			_.ReleaseNow();

			return AddNewBucket(BucketName);
		}
		else
		{
			return m_Buckets[It->second].get();
		}
	}

	BucketTracker* AddNewBucket(const std::string& BucketName)
	{
		RwLock::ExclusiveLockScope _(m_Lock);

		if (auto It = m_BucketMapping.find(BucketName); It == m_BucketMapping.end())
		{
			const uint32_t BucketIndex = gsl::narrow<uint32_t>(m_Buckets.size());
			m_Buckets.emplace_back(std::make_unique<BucketTracker>());
			m_BucketMapping[BucketName] = BucketIndex;

			return m_Buckets[BucketIndex].get();
		}
		else
		{
			return m_Buckets[It->second].get();
		}
	}

	RwLock										m_Lock;
	std::vector<std::unique_ptr<BucketTracker>> m_Buckets;
	tsl::robin_map<std::string, uint32_t>		m_BucketMapping;
};

struct ZenCacheTracker::Impl
{
	Impl(std::filesystem::path StateDirectory)
	{
		std::filesystem::path StatsDbPath{StateDirectory / ".zdb"};

		std::string RocksdbPath = ToUtf8(StatsDbPath);

		ZEN_DEBUG("opening tracker db at '{}'", RocksdbPath);

		rocksdb::DB*	   Db = nullptr;
		rocksdb::DBOptions Options;
		Options.create_if_missing = true;

		std::vector<std::string> ExistingColumnFamilies;
		rocksdb::Status			 Status = rocksdb::DB::ListColumnFamilies(Options, RocksdbPath, &ExistingColumnFamilies);

		std::vector<rocksdb::ColumnFamilyDescriptor> ColumnDescriptors;

		if (Status.IsPathNotFound())
		{
			ColumnDescriptors.emplace_back(rocksdb::ColumnFamilyDescriptor{rocksdb::kDefaultColumnFamilyName, {}});
		}
		else if (Status.ok())
		{
			for (const std::string& Column : ExistingColumnFamilies)
			{
				rocksdb::ColumnFamilyDescriptor ColumnFamily;
				ColumnFamily.name = Column;
				ColumnDescriptors.push_back(ColumnFamily);
			}
		}
		else
		{
			throw std::runtime_error("column family iteration failed for '{}': '{}'"_format(RocksdbPath, Status.getState()).c_str());
		}

		Status = rocksdb::DB::Open(Options, RocksdbPath, ColumnDescriptors, &m_RocksDbColumnHandles, &Db);

		if (!Status.ok())
		{
			throw std::runtime_error("database open failed for '{}': '{}'"_format(RocksdbPath, Status.getState()).c_str());
		}

		m_RocksDb.reset(Db);
	}

	~Impl()
	{
		for (auto* Column : m_RocksDbColumnHandles)
		{
			delete Column;
		}

		m_RocksDbColumnHandles.clear();
	}

	struct KeyStruct
	{
		uint16_t BucketId;
		IoHash	 HashKey;
	};

	struct ValueStruct
	{
		uint32_t CreateTime	 = 0;
		uint32_t AccessTime	 = 0;
		uint32_t AccessCount = 0;
		uint32_t LastGcCount = 0;
	};

	void TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
	{
		const uint32_t		  Ts = GetCurrentTimeStamp();
		rocksdb::WriteOptions Wo;
		rocksdb::ReadOptions  Ro;

		ZEN_UNUSED(BucketSegment);

		ValueStruct Value;

		rocksdb::Slice ValueSlice{(char*)&Value, sizeof Value};

		KeyStruct Key;
		Key.BucketId = 0;
		Key.HashKey	 = HashKey;

		rocksdb::Slice KeySlice{(char*)&Key, sizeof Key};

		std::string		ValueString;
		rocksdb::Status Status = m_RocksDb->Get(Ro, KeySlice, &ValueString);

		if (Status.ok() && ValueString.size() == sizeof(ValueStruct))
		{
			memcpy(&Value, ValueString.data(), ValueString.size());

			++Value.AccessCount;
		}
		else
		{
			Value.CreateTime  = Ts;
			Value.AccessCount = 1;
		}

		Value.AccessTime = Ts;

		m_RocksDb->Put(Wo, KeySlice, ValueSlice);
	}

	std::unique_ptr<rocksdb::DB>			  m_RocksDb;
	std::vector<rocksdb::ColumnFamilyHandle*> m_RocksDbColumnHandles;
	CacheAccessSnapshot						  m_CurrentSnapshot;
};

ZenCacheTracker::ZenCacheTracker(std::filesystem::path StateDirectory) : m_Impl(new Impl(StateDirectory))
{
}

ZenCacheTracker::~ZenCacheTracker()
{
	delete m_Impl;
}

void
ZenCacheTracker::TrackAccess(std::string_view BucketSegment, const IoHash& HashKey)
{
	m_Impl->TrackAccess(BucketSegment, HashKey);
}

#if ZEN_WITH_TESTS

TEST_CASE("z$.tracker")
{
	using namespace fmt::literals;
	using namespace std::literals;

	ScopedTemporaryDirectory TempDir;

	ZenCacheTracker Zcs(TempDir.Path());

	for (int i = 0; i < 10000; ++i)
	{
		IoHash KeyHash = IoHash::HashBuffer(&i, sizeof i);

		Zcs.TrackAccess("foo", KeyHash);
	}

	for (int i = 0; i < 10000; ++i)
	{
		IoHash KeyHash = IoHash::HashBuffer(&i, sizeof i);

		Zcs.TrackAccess("foo", KeyHash);
	}
}

#endif

void
cachetracker_forcelink()
{
}

}  // namespace zen