#pragma once

#include "stdint.h"
#include "cassert"
#include <atomic>
#include "Util/Span.h"

template <typename T>
inline T const* ByteToData(uint8_t const* data) { return reinterpret_cast<T const*>(data); }

template <typename T>
inline T* ByteToData(uint8_t* data) { return reinterpret_cast<T*>(data); }

template <typename T, typename MetaData>
class SharedBuffer final
{
private:
	struct ControlBlock
	{
		ControlBlock(uint32_t size, const MetaData& data) : Size{ size }, Data { data } {}

		std::atomic<uint32_t> Count{ 1 };
		uint32_t Size{};
		MetaData Data{};
	};

	static constexpr size_t ControlBlockOffset = 0;
	static constexpr size_t DataOffset = sizeof(ControlBlock);

public:
	SharedBuffer() = default;
	SharedBuffer(int size, const MetaData& metaData)
	{
		Data = new uint8_t[sizeof(ControlBlock) + size * sizeof(T)]{};

		new (GetControlBlock()) ControlBlock(size, metaData);
		for (size_t i{}; i < size; ++i)
		{
			new (&Ptr()[i]) T{};
		}
	}

	~SharedBuffer()
	{
		Destruct();
	}

	SharedBuffer(const SharedBuffer& other)
		: Data{ other.Data }
	{
		if (Data) ++GetControlBlock()->Count;
	}
	SharedBuffer(SharedBuffer&& other) noexcept
		: Data{ other.Data }
	{
		other.Data = nullptr;
	}
	
	SharedBuffer& operator=(const SharedBuffer& other)
	{
		if (this != &other) 
		{  
			Destruct();
			Data = other.Data;
			if (Data) ++GetControlBlock()->Count;
		}
    	return *this;
	}
	SharedBuffer& operator=(SharedBuffer&& other) noexcept
	{
		if (this != &other) 
		{
			Destruct();
			Data = other.Data;
			other.Data = nullptr;
		}
		return *this;
	}


private:
	ControlBlock const* GetControlBlock() const { assert(Data); return ByteToData<ControlBlock>(Data + ControlBlockOffset); }
	ControlBlock* GetControlBlock() { assert(Data); return ByteToData<ControlBlock>(Data + ControlBlockOffset); }

	void Destruct()
	{
		if (Data != nullptr && --GetControlBlock()->Count == 0)
		{
			uint32_t size = GetSize();
			GetControlBlock()->~ControlBlock();
			for (size_t i = 0; i < size; ++i)
				Ptr()[i].~T();
			delete[] Data;
		}
	}

public:
	T const* Ptr() const { assert(Data); return ByteToData<T>(Data + DataOffset); }
	T* Ptr() { assert(Data); return ByteToData<T>(Data + DataOffset); }
	tcb::span<const T> GetData() const { return tcb::span<const T>(Ptr(), GetSize()); }
	tcb::span<T> GetData() { return tcb::span<T>(Ptr(), GetSize()); }
	MetaData const* GetMetaData() const { assert(Data); return &GetControlBlock()->Data; }
	MetaData* GetMetaData() { assert(Data); return &GetControlBlock()->Data; }

	T& operator[] (uint32_t index) { assert(index < GetControlBlock()->Size); return GetData()[index]; }
	const T& operator[] (uint32_t index) const { assert(index < GetControlBlock()->Size); return GetData()[index]; }
	uint32_t GetSize() const { return GetControlBlock()->Size; }

	operator bool() const { return Data; }

private:
	uint8_t* Data{};
};