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

#pragma once

#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/logging.h>
#include <zencore/memory.h>
#include <zencore/thread.h>
#include <zencore/uid.h>
#include <zencore/zencore.h>

#pragma warning(push)
#pragma warning(disable : 4127)
#include <tsl/robin_map.h>
#pragma warning(pop)

#include <asio.hpp>

#include <chrono>

struct ZenCacheValue;

namespace spdlog {
class logger;
}

namespace zen {

class CbObjectWriter;
class ZenStructuredCacheClient;

/** Zen mesh tracker
 *
 * Discovers and tracks local peers
 */
class Mesh
{
public:
	Mesh(asio::io_context& IoContext);
	~Mesh();

	void Start(uint16_t Port);
	void Stop();

private:
	void Run();
	void IssueReceive();
	void EnqueueTick();
	void OnTick();
	void BroadcastPacket(CbObjectWriter&);

	enum State
	{
		kInitializing,
		kRunning,
		kExiting
	};

	static const int kMaxMessageSize = 2048;
	static const int kMaxUpdateSize	 = 1400;  // We'll try not to send messages larger than this

	spdlog::logger& Log() { return m_Log; }

	spdlog::logger&						   m_Log;
	std::atomic<State>					   m_State = kInitializing;
	asio::io_context&					   m_IoContext;
	std::unique_ptr<asio::ip::udp::socket> m_UdpSocket;
	std::unique_ptr<asio::ip::udp::socket> m_BroadcastSocket;
	asio::ip::udp::endpoint				   m_SenderEndpoint;
	std::unique_ptr<std::thread>		   m_Thread;
	uint16_t							   m_Port = 0;
	uint8_t								   m_MessageBuffer[kMaxMessageSize];
	asio::high_resolution_timer			   m_Timer{m_IoContext};
	Oid									   m_SessionId;

	struct PeerInfo
	{
		Oid							   SessionId;
		std::time_t					   LastSeen;
		std::vector<asio::ip::address> SeenOnIP;
	};

	RwLock									   m_SessionsLock;
	tsl::robin_map<Oid, PeerInfo, Oid::Hasher> m_KnownPeers;
};

namespace detail {
	struct ZenCacheSessionState;
}

struct ZenCacheResult
{
	IoBuffer	Response;
	int64_t		Bytes		   = {};
	double		ElapsedSeconds = {};
	int32_t		ErrorCode	   = {};
	std::string Reason;
	bool		Success = false;
};

/** Zen Structured Cache session
 *
 * This provides a context in which cache queries can be performed
 *
 * These are currently all synchronous. Will need to be made asynchronous
 */
class ZenStructuredCacheSession
{
public:
	ZenStructuredCacheSession(ZenStructuredCacheClient& OuterClient);
	~ZenStructuredCacheSession();

	ZenCacheResult CheckHealth();
	ZenCacheResult GetCacheRecord(std::string_view BucketId, const IoHash& Key, ZenContentType Type);
	ZenCacheResult GetCachePayload(std::string_view BucketId, const IoHash& Key, const IoHash& PayloadId);
	ZenCacheResult PutCacheRecord(std::string_view BucketId, const IoHash& Key, IoBuffer Value, ZenContentType Type);
	ZenCacheResult PutCachePayload(std::string_view BucketId, const IoHash& Key, const IoHash& PayloadId, IoBuffer Payload);

private:
	inline spdlog::logger& Log() { return m_Log; }

	spdlog::logger&				  m_Log;
	ZenStructuredCacheClient&	  m_Client;
	detail::ZenCacheSessionState* m_SessionState;
};

/** Zen Structured Cache client
 *
 * This represents an endpoint to query -- actual queries should be done via
 * ZenStructuredCacheSession
 */
class ZenStructuredCacheClient : public RefCounted
{
public:
	ZenStructuredCacheClient(std::string_view ServiceUrl);
	~ZenStructuredCacheClient();

	std::string_view ServiceUrl() const { return m_ServiceUrl; }

	inline spdlog::logger& Log() { return m_Log; }

private:
	spdlog::logger& m_Log;
	std::string		m_ServiceUrl;

	RwLock									 m_SessionStateLock;
	std::list<detail::ZenCacheSessionState*> m_SessionStateCache;

	detail::ZenCacheSessionState* AllocSessionState();
	void						  FreeSessionState(detail::ZenCacheSessionState*);

	friend class ZenStructuredCacheSession;
};

}  // namespace zen