#pragma once

#include <sp/common/DataBuffer.h>
#include <sp/protocol/BitField.h>

namespace sp {

// TODO: flush if offset exceeds 64
class BitBuffer {
  private:
	using Data = std::uint64_t;

  private:
	DataBuffer& m_Buffer;
	Data m_Data;
	std::size_t m_Offset;
	bool m_WasBitField;

  public:
	BitBuffer(DataBuffer& a_Buffer) : m_Buffer(a_Buffer), m_Data(0), m_Offset(0), m_WasBitField(false) {}

	void UpdateWrite(bool a_IsBitField) {
		if (m_Offset == 0)
			return;
		if ((m_Offset % 8 == 0) || (!a_IsBitField && m_WasBitField)) {
			Flush();
		}
		m_WasBitField = a_IsBitField;
	}

	void UpdateRead(bool a_IsBitField) {
		if (m_Offset == 0)
			return;
		if ((m_Offset % 8 == 0) || (!a_IsBitField && m_WasBitField)) {
			MoveReadOffset();
		}
		m_WasBitField = a_IsBitField;
	}

	template <typename T, std::size_t BitSize>
	void Append(T a_Data) {
		Data bin = static_cast<Data>(a_Data);
		bin &= ((1 << BitSize) - 1);  // prevents overflow
		std::size_t pushCount = sizeof(Data) * 8 - m_Offset - BitSize;
		m_Data |= bin << pushCount;
		m_Offset += BitSize;
	}

	template <typename T, std::size_t BitSize>
	void Read(T& a_Data) {
		std::size_t byteCount = GetByteCount(m_Offset + BitSize);
		constexpr Data dataMask = (1 << BitSize) - 1;
		m_Buffer.ReadSome(reinterpret_cast<std::uint8_t*>(&m_Data), byteCount);
		SwapBytes(reinterpret_cast<std::uint8_t*>(&m_Data), reinterpret_cast<std::uint8_t*>(&m_Data) + byteCount);
		m_Data >>= byteCount * 8 - m_Offset - BitSize;
		m_Data &= dataMask;
		a_Data = T(m_Data);
		m_Buffer.SetReadOffset(m_Buffer.GetReadOffset() - byteCount);
		m_Offset += BitSize;
	}

  private:
	void Flush() {
		std::size_t byteCount = GetByteCount();
		m_Data >>= (sizeof(Data) - byteCount) * 8;
		SwapBytes(reinterpret_cast<std::uint8_t*>(&m_Data), reinterpret_cast<std::uint8_t*>(&m_Data) + byteCount);
		m_Buffer.WriteSome(reinterpret_cast<std::uint8_t*>(&m_Data), byteCount);
		m_Offset = 0;
		m_WasBitField = false;
		m_Data = 0;
	}

	void MoveReadOffset() {
		std::size_t byteCount = GetByteCount();
		m_Buffer.SetReadOffset(m_Buffer.GetReadOffset() + byteCount);
		m_Offset = 0;
		m_WasBitField = false;
		m_Data = 0;
	}

	std::size_t GetByteCount(std::size_t a_Offset = -1) const {
		if (a_Offset == static_cast<std::size_t>(-1))
			a_Offset = m_Offset;
		if (a_Offset <= 8)
			return 1;
		return (a_Offset - 1) / 8 + 1;
	}
};

}  // namespace sp