#pragma once

#include <sp/common/VarInt.h>
#include <sp/protocol/MessageFactory.h>

namespace sp {

template <typename TMessageFactory>
DataBuffer MessageStream<TMessageFactory>::ReadAndDecapsulate() {
	VarInt messageLength;
	messageLength.Read([this](std::uint8_t& data) {
		DataBuffer buffer = m_Stream->Read(1);
		data = *buffer.data();
	});
	std::size_t amount = messageLength.GetValue();
	DataBuffer buffer = m_Stream->Read(amount);
	
	for (auto& enc : m_Encapsulators) {
		buffer = enc->Decapsulate(buffer);
	}

	return buffer;
}

template <typename TMessageFactory>
std::unique_ptr<typename TMessageFactory::MessageBaseType> MessageStream<TMessageFactory>::MakeMessage(DataBuffer& buffer, MessageIdType a_Id) {
	static const TMessageFactory FACTORY;
	auto message = FACTORY.CreateMessage(a_Id);
	message->Read(buffer);
	return message;
}

template <typename TMessageFactory>
std::unique_ptr<typename TMessageFactory::MessageBaseType> MessageStream<TMessageFactory>::ReadMessage(MessageIdType a_Id) {
	
	DataBuffer buffer = ReadAndDecapsulate();

	return MakeMessage(buffer, a_Id);
}

template <typename TMessageFactory>
std::unique_ptr<typename TMessageFactory::MessageBaseType> MessageStream<TMessageFactory>::ReadMessage() {
	DataBuffer buffer = ReadAndDecapsulate();

	VarInt messageId;
	buffer >> messageId;

	return MakeMessage(buffer, MessageIdType(messageId.GetValue()));
}

template <typename TMessageFactory>
void MessageStream<TMessageFactory>::WriteMessage(const MessageBaseType& a_Message, bool a_WriteId) {
	DataBuffer buffer;
	if (a_WriteId)
		buffer << VarInt{static_cast<std::uint64_t>(a_Message.GetId())};
	buffer << a_Message.Write();
	for (auto& enc : m_Encapsulators) {
		buffer = enc->Encapsulate(buffer);
	}
	DataBuffer header;
	header << VarInt{buffer.GetSize()};
	m_Stream->Write(header);
	m_Stream->Write(buffer);
}

}  // namespace sp