// 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>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
#include <asio.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

#include <chrono>
#include <list>

struct ZenCacheValue;

namespace spdlog {
class logger;
}

namespace zen {

class CbObjectWriter;
class CbObjectView;
class CbPackage;
class ZenStructuredCacheClient;

/** Zen mesh tracker
 *
 * Discovers and tracks local peers
 *
 * NOTE: This is currently experimental, and not very useful yet
 *
 */

class MeshTracker
{
public:
	MeshTracker(asio::io_context& IoContext);
	~MeshTracker();

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

struct ZenStructuredCacheClientOptions
{
	std::string_view			 Name;
	std::string_view			 Url;
	std::span<std::string const> Urls;
	std::chrono::milliseconds	 ConnectTimeout{};
	std::chrono::milliseconds	 Timeout{};
};

/** 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 GetCacheValue(std::string_view BucketId, const IoHash& Key, const IoHash& ValueContentId);
	ZenCacheResult PutCacheRecord(std::string_view BucketId, const IoHash& Key, IoBuffer Value, ZenContentType Type);
	ZenCacheResult PutCacheValue(std::string_view BucketId, const IoHash& Key, const IoHash& ValueContentId, IoBuffer Payload);
	ZenCacheResult InvokeRpc(const CbObjectView& Request);
	ZenCacheResult InvokeRpc(const CbPackage& Package);

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(const ZenStructuredCacheClientOptions& Options);
	~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;
	std::chrono::milliseconds m_ConnectTimeout;
	std::chrono::milliseconds m_Timeout;

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

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

	friend class ZenStructuredCacheSession;
};

}  // namespace zen