/*
 * 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 Stable Vector
	 * @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, size_t GROW_FACTOR = 2, size_t DEFAULT_CHUNK_SIZE = 32> 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_size; 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++;
					}
				}

				return reqIndex + gapCount;
			}

			T& GetAlignedData(size_t index)
			{
				size_t realIndex = GetRealIndex(index);

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

			size_t GetLastOccupiedIndex()
			{
				for (size_t i = m_capacity - 1; i >= 0; i--)
				{
					if (m_occupiedIndices[i])
					{
						return i;
					}
				}

				return 0;
			}
		};

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

			T& operator*() { return m_chunk->GetAlignedData(m_index); }
			T* operator->() { return &m_chunk->GetAlignedData(m_index); }

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

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

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

			Iterator operator--(int)
			{
				Iterator temp = *this;
				--m_index;
				MoveToPrevChunk();
				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 m_chunk != other.m_chunk || m_index != other.m_index;
			}

		protected:
			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() {}

		private:
			VectorChunk* m_chunk;
			size_t m_index;
		};

	public:
		StableVector() : m_FirstChunk(nullptr), m_LastChunk(nullptr), m_Size(0), m_Capacity(0)
		{
			m_FirstChunk = Grow(DEFAULT_CHUNK_SIZE);
			m_LastChunk = m_FirstChunk;
			m_Capacity = DEFAULT_CHUNK_SIZE;
		}

		StableVector(size_t 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)
		    : m_FirstChunk(nullptr), m_LastChunk(nullptr), m_Size(0), m_Capacity(0)
		{
			m_FirstChunk = Grow(size);
			m_LastChunk = m_FirstChunk;
			m_Capacity = size;

			for (size_t i = 0; i < size; i++)
			{
				PushBack(value);
			}
		}

		StableVector(std::initializer_list<T> list)
		    : m_FirstChunk(nullptr), m_LastChunk(nullptr), m_Size(0), m_Capacity(0)
		{
			m_FirstChunk = Grow(list.size());
			m_LastChunk = m_FirstChunk;
			m_Capacity = list.size();

			for (const T& value: list)
			{
				PushBack(value);
			}
		}

		StableVector(const StableVector<T>& copy)
		    : m_FirstChunk(nullptr), m_LastChunk(nullptr), m_Size(0), m_Capacity(0)
		{
			m_FirstChunk = Grow(copy.m_Capacity); // One big chunk to make Stable contiguous.
			m_LastChunk = m_FirstChunk;
			m_Capacity = copy.m_Capacity;

			for (size_t i = 0; i < copy.Size(); i++)
			{
				PushBack(copy.At(i));
			}
		}

		StableVector(StableVector<T>&& move) noexcept
		    : m_FirstChunk(nullptr), m_LastChunk(nullptr), m_Size(0), m_Capacity(0)
		{
			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()
		{
			VectorChunk* currentChunk = m_FirstChunk;
			while (currentChunk)
			{
				VectorChunk* nextChunk = currentChunk->m_next;
				currentChunk->~VectorChunk();
				::operator delete(currentChunk);
				currentChunk = nextChunk;
			}
		}

		void PushBack(const T& value)
		{
			if (m_LastChunk->m_nextIndex == m_LastChunk->m_capacity)
			{
				VectorChunk* newChunk = Grow(m_LastChunk->m_capacity * GROW_FACTOR);
				m_LastChunk->m_next = newChunk;
				newChunk->m_prev = m_LastChunk;
				m_LastChunk = newChunk;
			}

			new (&m_LastChunk->m_data[m_LastChunk->m_nextIndex]) T(value);
			m_LastChunk->m_occupiedIndices[m_LastChunk->m_nextIndex] = true;
			m_LastChunk->m_nextIndex++;
			m_LastChunk->m_size++;
			m_Size++;
		}

		/// std version of push_back(const T& value)
		void push_back(const T& value) { PushBack(value); }

		void PushBack(T&& value) noexcept
		{
			if (m_LastChunk->m_nextIndex == m_LastChunk->m_capacity)
			{
				VectorChunk* newChunk = Grow(m_LastChunk->m_capacity * GROW_FACTOR);
				m_LastChunk->m_next = newChunk;
				newChunk->m_prev = m_LastChunk;
				m_LastChunk = newChunk;
			}

			new (&m_LastChunk->m_data[m_LastChunk->m_nextIndex]) T(std::move(value));
			m_LastChunk->m_occupiedIndices[m_LastChunk->m_nextIndex] = true;
			m_LastChunk->m_nextIndex++;
			m_LastChunk->m_size++;
			m_Size++;
		}

		/// std version of push_back(T&& value)
		void push_back(T&& value) { PushBack(value); }

		// Checks the first available gap and inserts. If no gap it works like push_back(const T& value)
		void Push(const T& value)
		{
			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])
						{
							new (&currentChunk->m_data[i]) T(value);
							currentChunk->m_occupiedIndices[i] = true;
							currentChunk->m_size++;
							currentChunk->m_gapCount--;
							m_Size++;
							return;
						}
					}
				}
			}

			PushBack(value);
		}

		// Checks the first available gap and inserts. If no gap it works like push_back(T&& value)
		void Push(T&& value) noexcept
		{
			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])
						{
							new (&currentChunk->m_data[i]) T(std::move(value));
							currentChunk->m_occupiedIndices[i] = true;
							currentChunk->m_size++;
							currentChunk->m_gapCount--;
							m_Size++;
							return;
						}
					}
				}
			}

			PushBack(std::move(value));
		}

		template<typename... Args> void EmplaceBack(Args&&... args)
		{
			if (m_LastChunk->m_nextIndex == m_LastChunk->m_capacity)
			{
				VectorChunk* newChunk = Grow(m_LastChunk->m_capacity * GROW_FACTOR);
				m_LastChunk->m_next = newChunk;
				newChunk->m_prev = m_LastChunk;
				m_LastChunk = newChunk;
			}

			new (&m_LastChunk->m_data[m_LastChunk->m_nextIndex]) T(std::forward<Args>(args)...);
			m_LastChunk->m_occupiedIndices[m_LastChunk->m_nextIndex] = true;
			m_LastChunk->m_nextIndex++;
			m_LastChunk->m_size++;
			m_Size++;
		}

		/// std version of emplace_back
		template<typename... Args> void emplace_back(Args&&... args) { EmplaceBack(std::forward<Args>(args)...); }

		// Checks the first available gap and inserts. If no gap it works like emplace_back(Args&&... args)
		template<typename... Args> void Emplace(Args&&... args)
		{
			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])
						{
							new (&currentChunk->m_data[i]) T(std::forward<Args>(args)...);
							currentChunk->m_occupiedIndices[i] = true;
							currentChunk->m_size++;
							currentChunk->m_gapCount--;
							m_Size++;
							return;
						}
					}
				}
			}

			EmplaceBack(std::forward<Args>(args)...);
		}

		void PopBack()
		{
			if (m_Size == 0)
			{
				return; // return? or make
			}

			m_LastChunk->m_data[m_LastChunk->m_nextIndex - 1].~T();
			m_LastChunk->m_occupiedIndices[m_LastChunk->m_nextIndex - 1] = false;
			m_LastChunk->m_nextIndex--;
			m_LastChunk->m_size--;
			m_Size--;
		}

		/// std version of pop_back
		void pop_back() { PopBack(); }

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

			return m_LastChunk->m_data[m_LastChunk->GetLastOccupiedIndex()];
		}

		/// std version of back
		constexpr T& back() const { return Back(); }

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

			return m_FirstChunk->m_data[0];
		}

		/// std version of front
		constexpr T& front() const { return Front(); }

		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");
			}

			handle.first->m_data[realIndex].~T();
			handle.first->m_occupiedIndices[realIndex] = false;
			handle.first->m_size--;
			handle.first->m_gapCount++;
			m_Size--;
		}

		// Temporary non correct solution. (DONT MIND THIS ERASE)
		void erase(Iterator it)
		{
			size_t index = 0;
			VectorChunk* currentChunk = m_FirstChunk;
			while (currentChunk)
			{
				if (index + currentChunk->m_size > it.m_index)
				{
					size_t realIndex = currentChunk->GetRealIndex(it.m_index);
					currentChunk->m_data[realIndex].~T();
					currentChunk->m_occupiedIndices[realIndex] = false;
					currentChunk->m_size--;
					currentChunk->m_gapCount++;
					m_Size--;
					return;
				}

				index += currentChunk->m_size;
				currentChunk = currentChunk->m_next;
			}
		}

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

			auto handle = FindChunk(index);
			return handle.first->GetAlignedData(handle.second);
		}

		/// std version of at
		T& at(size_t index) { return At(index); }

		size_t Size() const noexcept { return m_Size; }
		/// std version of size
		size_t size() const noexcept { return m_Size; }

		size_t Capacity() const noexcept { return m_Capacity; }
		/// std version of capacity
		size_t capacity() const noexcept { return m_Capacity; }

		bool Empty() const noexcept { return m_Size == 0; }
		/// std version of empty
		bool empty() const noexcept { return m_Size == 0; }

		Iterator begin() { return Iterator(m_FirstChunk, 0); }
		Iterator end() { return Iterator(m_LastChunk, m_LastChunk->m_nextIndex - 1); }

		const Iterator& cbegin() { return Iterator(m_FirstChunk, 0); }
		const Iterator& cend() { return Iterator(m_LastChunk, m_LastChunk->m_nextIndex - 1); }

	protected:
		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");
		}

	private:
		VectorChunk* m_FirstChunk;
		VectorChunk* m_LastChunk;
		size_t m_Size;
		size_t m_Capacity;
	};
}

#pragma warning(pop)