#pragma once

#include "Utils/IO/DataStream.hpp"

#include <tsl/robin_map.h>
#include <tsl/robin_set.h>
#include <type_traits>

namespace DataStreamAdapters {
template <class TKeyAdapter = void, class TValueAdapter = void>
struct TslRobinMap
{
	template <class TKey, class TValue>
	static void ReadFromDataStream(InputDataStream& s, tsl::robin_map<TKey, TValue>& map)
	{
		static_assert(std::is_default_constructible_v<TValue>);
		static_assert(std::is_move_constructible_v<TValue>);

		uint64_t size;
		s.Read(size);
		map.reserve(size);

		for (uint64_t i = 0; i < size; ++i) {
			TKey key;
			if constexpr (std::is_same_v<TKeyAdapter, void>) {
				s.ReadGeneric(key);
			} else {
				s.ReadObjectAdapted<TKeyAdapter>(key);
			}

			TValue value;
			if constexpr (std::is_same_v<TValueAdapter, void>) {
				s.ReadGeneric(value);
			} else {
				s.ReadObjectAdapted<TValueAdapter>(value);
			}

			map.insert(std::move(key), std::move(value));
		}
	}

	template <class TKey, class TValue>
	static void WriteToDataStream(OutputDataStream& s, const tsl::robin_map<TKey, TValue>& map)
	{
		s.Write((uint64_t)map.size());

		for (auto it = map.begin(); it != map.end(); ++it) {
			if constexpr (std::is_same_v<TKeyAdapter, void>) {
				s.WriteGeneric(it.key());
			} else {
				s.WriteObjectAdapted<TKeyAdapter>(it.key());
			}

			if constexpr (std::is_same_v<TValueAdapter, void>) {
				s.WriteGeneric(it.value());
			} else {
				s.WriteObjectAdapted<TValueAdapter>(it.value());
			}
		}
	}
};

template <class TAdapter = void>
struct TslRobinSet
{
	template <class TElement>
	static void ReadFromDataStream(InputDataStream& s, tsl::robin_set<TElement>& set)
	{
		static_assert(std::is_default_constructible_v<TElement>);
		static_assert(std::is_move_constructible_v<TElement>);

		uint64_t size;
		s.Read(size);
		set.reserve(size);

		for (uint64_t i = 0; i < size; ++i) {
			TElement element;
			if constexpr (std::is_same_v<TAdapter, void>) {
				s.ReadGeneric(element);
			} else {
				s.ReadObjectAdapted<TAdapter>(element);
			}

			set.insert(std::move(element));
		}
	}

	template <class TElement>
	static void WriteToDataStream(OutputDataStream& s, const tsl::array_set<TElement>& set)
	{
		s.Write((uint64_t)set.size());

		for (auto& element : set) {
			if constexpr (std::is_same_v<TAdapter, void>) {
				s.WriteGeneric(element);
			} else {
				s.WriteObjectAdapted<TAdapter>(element);
			}
		}
	}
};
} // namespace DataStreamAdapters