/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#pragma once

#include "Base/Wrapper.hpp"

#include <initializer_list>
#include <iterator>
#include <stdexcept>
#include <list>
#include <vector>
#include <algorithm>

#include <iostream>

#pragma warning(push)
#pragma warning(disable : 4200)
#pragma warning(disable : 6011)
#pragma warning(disable : 6386)

namespace OpenVulkano
{
	/**
	 * @class StableVector
	 * @brief Stable Vector is an alternative version for std::vector that provides chunk based memory allocation without re-aligning
	 * 
	 * @throw Please know that this vector creates array gaps when you remove an element.
	 */
	template<typename T> class StableVector
	{
		struct VectorChunk
		{
			VectorChunk* m_next = nullptr; // Next chunk
			VectorChunk* m_prev = nullptr; // Previous chunk

			size_t m_size; // Size of the chunk
			size_t m_capacity; // Capacity of the chunk
			size_t m_nextIndex; // Next index to insert
			size_t m_gapCount; // Count of emptied gaps in the chunk
			bool* m_occupiedIndices; // filled gaps array
			T m_data[0]; // data array

			VectorChunk(size_t size) : m_size(0), m_capacity(size), m_nextIndex(0), m_gapCount(0)
			{
				m_occupiedIndices = reinterpret_cast<bool*>(m_data + size);
				memset(m_occupiedIndices, 0, size * sizeof(bool));
			}

			~VectorChunk()
			{
				for (size_t i = 0; i < m_capacity; i++)
				{
					if (m_occupiedIndices[i])
					{
						m_data[i].~T();
					}
				}
				m_prev = nullptr;
				m_next = nullptr;
			}

			size_t GetRealIndex(size_t reqIndex)
			{
				if (m_gapCount == 0)
				{
					return reqIndex;
				}

				size_t gapCount = 0;
				for (size_t i = 0; i <= reqIndex; i++)
				{
					if (!m_occupiedIndices[i])
					{
						gapCount++;
					}
				}

				for (size_t i = reqIndex + gapCount; i < m_capacity; i++)
				{
					if (m_occupiedIndices[i])
					{
						return i;
					}
				}

				return reqIndex + gapCount;
			}
		};

	public:
		template<typename U>
		class Iterator
		{
			friend class StableVector<T>;

		public:
			Iterator(VectorChunk* chunk, size_t index) : m_chunk(chunk), m_index(index) {}

			U& operator*()
			{
				auto realIndex = m_chunk->GetRealIndex(m_index);
				return m_chunk->m_data[realIndex];
			}

			U* operator->() { return &operator*(); }

			Iterator& operator++()
			{
				++m_index;
				MoveToNextChunk();
				return *this;
			}

			Iterator operator++(int)
			{
				Iterator temp = *this;
				++temp;
				return temp;
			}

			Iterator& operator--()
			{
				--m_index;
				MoveToPrevChunk();
				return *this;
			}

			Iterator operator--(int)
			{
				Iterator temp = *this;
				--temp;
				return temp;
			}

			Iterator operator+(size_t n)
			{
				Iterator temp = *this;
				temp.m_index += n;
				temp.MoveToNextChunk();
				return temp;
			}

			Iterator operator-(size_t n)
			{
				Iterator temp = *this;
				temp.m_index -= n;
				temp.MoveToPrevChunk();
				return temp;
			}

			bool operator==(const Iterator& other) const
			{
				return m_chunk == other.m_chunk && m_index == other.m_index;
			}

			bool operator!=(const Iterator& other) const { return !(*this == other); }

		private:
			void MoveToNextChunk()
			{
				while (m_chunk && m_index >= m_chunk->m_size)
				{
					m_index -= m_chunk->m_size;
					m_chunk = m_chunk->m_next;
				}
			}

			void MoveToPrevChunk()
			{
				while (m_chunk && m_index < 0)
				{
					m_index += m_chunk->m_size;
					m_chunk = m_chunk->m_prev;
				}
			}

			size_t GetIteratorIndex() const { return m_chunk->GetRealIndex(m_index); }

		private:
			VectorChunk* m_chunk;
			size_t m_index;
		};

	public:
		using ConstIterator = Iterator<const T>;
		using RegIterator = Iterator<T>;

		StableVector(size_t size = DEFAULT_CHUNK_SIZE)
		    : m_firstChunk(nullptr), m_lastChunk(nullptr), m_size(0), m_capacity(0)
		{
			m_firstChunk = Grow(size);
			m_lastChunk = m_firstChunk;
			m_capacity = size;
		}

		StableVector(size_t size, const T& value) : StableVector(size)
		{
			for (size_t i = 0; i < size; i++)
			{
				PushBack(value);
			}
		}

		StableVector(std::initializer_list<T> list) : StableVector(list.size)
		{
			for (const T& value: list)
			{
				PushBack(value);
			}
		}

		StableVector(const StableVector<T>& copy) : StableVector(copy.m_size)
		{
			for (size_t i = 0; i < copy.Size(); i++)
			{
				PushBack(copy.At(i));
			}
		}

		StableVector(StableVector<T>&& move) noexcept : m_firstChunk(nullptr)
		{
			m_firstChunk = move.m_firstChunk;
			m_lastChunk = move.m_lastChunk;
			m_size = move.m_size;
			m_capacity = move.m_capacity;

			move.m_firstChunk = nullptr;
			move.m_lastChunk = nullptr;
			move.m_size = 0;
			move.m_capacity = 0;
		}

		~StableVector()
		{
			// Destroy everything left
			VectorChunk* currentChunk = m_firstChunk;
			while (currentChunk)
			{
				VectorChunk* temp = currentChunk;
				currentChunk = currentChunk->m_next;
				temp->~VectorChunk();
				::operator delete(temp);
			}
		}

		void PushBack(const T& value) { new (InsertBack()) T(value); }
		void PushBack(T&& value) noexcept { new (InsertBack()) T(std::move(value)); }

		template<typename... Args> void EmplaceBack(Args&&... args)
		{
			new (InsertBack()) T(std::forward<Args>(args)...);
		}

		void Push(const T& value) { new (Insert()) T(value); }
		void Push(T&& value) noexcept { new (Insert()) T(std::move(value)); }
		template<typename... Args> void Emplace(Args&&... args) { new (Insert()) T(std::forward<Args>(args)...); }

		void PopBack()
		{
			if (m_size == 0) return;

			if (m_lastChunk->m_nextIndex == 1 && m_lastChunk != m_firstChunk)
			{
				VectorChunk* temp = m_lastChunk;
				m_capacity -= temp->m_capacity;
				m_lastChunk = m_lastChunk->m_prev;
				m_lastChunk->m_next = nullptr;
				temp->~VectorChunk();
				::operator delete(temp);
				m_size--;
				return;
			}

			m_lastChunk->m_nextIndex--;
			m_lastChunk->m_data[m_lastChunk->m_nextIndex].~T();
			m_lastChunk->m_occupiedIndices[m_lastChunk->m_nextIndex] = false;
			m_lastChunk->m_size--;
			m_lastChunk->m_gapCount++;
			m_size--;
		}

		constexpr T& Back() const
		{
			if (m_size == 0)
			{
				throw std::out_of_range("Index out of range");
			}

			return m_lastChunk->m_data[m_lastChunk->m_nextIndex - 1];
		}

		constexpr T& Front() const
		{
			if (m_size == 0)
			{
				throw std::out_of_range("Index out of range");
			}

			return m_firstChunk->m_data[0];
		}

		void Remove(size_t index)
		{
			auto handle = FindChunk(index);
			if (!handle.first)
			{
				throw std::out_of_range("Index out of range");
			}

			size_t realIndex = handle.first->GetRealIndex(handle.second);
			if (realIndex >= handle.first->m_size)
			{
				throw std::out_of_range("Index out of range");
			}

			Delete(realIndex, handle.first);
		}

		void Remove(const RegIterator& it)
		{
			if (it.m_chunk == nullptr)
			{
				throw std::out_of_range("Index out of range");
			}

			Delete(it.GetIteratorIndex(), it.m_chunk);
		}

		T& operator[](size_t index) noexcept
		{
			auto handle = FindChunk(index);
			size_t realIndex = handle.first->GetRealIndex(handle.second);
			return handle.first->m_data[realIndex];
		}

		T& At(size_t index)
		{
			if (index >= m_size) [[unlikely]]
			{
				throw std::out_of_range("Index out of range");
			}

			return operator[](index);
		}

		void Clear()
		{
			// Get the first chunk's cap
			size_t cap = m_firstChunk->m_capacity;

			// Annihilate everything
			VectorChunk* currentChunk = m_firstChunk->m_next;
			while (currentChunk)
			{
				VectorChunk* temp = currentChunk;
				currentChunk = currentChunk->m_next;
				temp->~VectorChunk();
				::operator delete(temp);
			}

			// Reset the first chunk
			m_firstChunk->m_next = nullptr;
			m_lastChunk = m_firstChunk;
			m_firstChunk->m_size = 0;
			m_firstChunk->m_nextIndex = 0;
			m_firstChunk->m_gapCount = 0;
			m_firstChunk->~VectorChunk();
			new (m_firstChunk) VectorChunk(cap);
			m_size = 0;
			m_capacity = cap;
		}

		size_t Size() const noexcept { return m_size; }
		size_t Capacity() const noexcept { return m_capacity; }
		bool Empty() const noexcept { return m_size == 0; }

		RegIterator begin() { return RegIterator(m_firstChunk, 0); }
		RegIterator end()
		{ 
			if (Empty()) [[unlikely]]
			{
				return begin();
			}
			return RegIterator(m_lastChunk, m_lastChunk->m_nextIndex - 1); 
		}

		ConstIterator cbegin() const { return ConstIterator(m_firstChunk, 0); }
		ConstIterator cend() const
		{
			if (Empty()) [[unlikely]]
			{
				return cbegin();
			}
			return ConstIterator(m_lastChunk, m_lastChunk->m_nextIndex - 1); 
		}

		//region std aliases
		void push_back(const T& value) { PushBack(value); }
		void push_back(T&& value) { PushBack(std::move(value)); }
		template<typename... Args> void emplace_back(Args&&... args) { EmplaceBack(std::forward<Args>(args)...); }
		void pop_back() { PopBack(); }
		T& back() { return Back(); }
		T& front() { return Front(); }
		bool empty() const noexcept { return Empty(); }
		void clear() { Clear(); }
		void erase(const RegIterator& it) { Remove(it); }
		size_t size() const noexcept { return Size(); }
		size_t capacity() const noexcept { return Capacity(); }
		//endregion

	private:
		VectorChunk* Grow(size_t requestSize)
		{
			VectorChunk* newChunk = static_cast<VectorChunk*>(
			::operator new(sizeof(VectorChunk) + requestSize * sizeof(T) + requestSize * sizeof(bool)));
			new (newChunk) VectorChunk(requestSize);

			return newChunk;
		}

		std::pair<VectorChunk*, size_t> FindChunk(size_t index)
		{
			size_t leftIndex = index;
			VectorChunk* currentChunk = m_firstChunk;
			while (currentChunk)
			{
				if (leftIndex < currentChunk->m_size)
				{
					return { currentChunk, leftIndex };
				}

				leftIndex -= currentChunk->m_size;
				currentChunk = currentChunk->m_next;
			}

			throw std::out_of_range("Index out of range");
		}

		T* InsertBack()
		{
			if (m_lastChunk->m_nextIndex == m_lastChunk->m_capacity)
			{
				VectorChunk* newChunk = Grow(m_lastChunk->m_capacity * GROW_FACTOR);
				m_capacity += m_lastChunk->m_capacity * GROW_FACTOR;
				m_lastChunk->m_next = newChunk;
				newChunk->m_prev = m_lastChunk;
				m_lastChunk = newChunk;
			}

			m_lastChunk->m_occupiedIndices[m_lastChunk->m_nextIndex] = true;
			m_lastChunk->m_size++;
			m_lastChunk->m_gapCount++;
			m_size++;
			return &m_lastChunk->m_data[m_lastChunk->m_nextIndex++];
		}

		T* Insert()
		{
			T* target = nullptr;
			VectorChunk* currentChunk = m_firstChunk;
			while (currentChunk)
			{
				if (currentChunk->m_gapCount > 0) // If there is a gap check occupied indices to find the first gap
				{
					for (size_t i = 0; i < currentChunk->m_capacity; i++)
					{
						if (!currentChunk->m_occupiedIndices[i])
						{
							target = &currentChunk->m_data[i];
							currentChunk->m_occupiedIndices[i] = true;
							currentChunk->m_size++;
							currentChunk->m_gapCount--;
							m_size++;
							return target;
						}
					}
				}

				currentChunk = currentChunk->m_next;
			}

			return InsertBack();
		}

		void Delete(size_t index, VectorChunk* chunk)
		{
			chunk->m_data[index].~T();
			chunk->m_occupiedIndices[index] = false;
			chunk->m_size--;
			chunk->m_gapCount++;
			m_size--;
		}

	private:
		VectorChunk* m_firstChunk;
		VectorChunk* m_lastChunk;
		size_t m_size;
		size_t m_capacity;

		static inline constexpr size_t DEFAULT_CHUNK_SIZE = 32;
		static inline constexpr size_t GROW_FACTOR = 2;
	};
}

#pragma warning(pop)