// License: Modified MIT (NON-AI)
// See the LICENSE file in the root directory for license information.
// Copyright 2025 Timothy Prepscius

#pragma once

#include "StackArray.hpp"
#include "../exception/OutOfBounds.hpp"
#include "../algorithm/vector_erase_if_value.hpp"

namespace tjp::core {

template<typename T>
using MultiContainer_ = StackArray<T *>;

template<typename T, typename I, typename J>
struct MultiIterator
{
	using Self = MultiIterator<T,I,J>;
	using value_type = typename T::value_type;
	using reference = typename T::reference;
	
	I end;
	I i;
	J j;
	bool ended;

	MultiIterator()
	{
	}

	MultiIterator(I &&end_, I &&i_, J &&j_, bool ended_) :
		end(end_),
		i(i_),
		j(j_),
		ended(ended_)
	{
		validate();
	}
	
	auto &deref ()
	{
		return *j;
	}
	
	auto &operator *() { return deref(); }
	auto &operator ->() { return deref(); }
	
	void validate()
	{
		while (!ended && j == (*i)->end())
		{
			++i;
			
			if (i == end)
				ended = true;
			else
				j = (*i)->begin();
		}
	}
	
	void next()
	{
		if (ended)
			return;

		++j;
		validate();
	}
	
	auto &operator++(int)
	{
		next();
		return *this;
	}

	auto operator++()
	{
		auto v = *this;
		next();
		return v;
	}

	bool operator ==(const Self &rhs) const
	{
		return
			(ended && rhs.ended) ||
			(i == rhs.i && j == rhs.j);
	}
	
	bool operator !=(const Self &rhs) const
	{
		return !(*this == rhs);
	}
} ;

template<typename T, typename C=MultiContainer_<T>>
struct MultiContainer {
	C c;
	
	using Self = MultiContainer<T, C>;
	using iterator = MultiIterator<T,typename C::iterator,typename T::iterator>;
	using const_iterator = MultiIterator<T,typename C::const_iterator,typename T::const_iterator>;
	using value_type = typename T::value_type;
	using reference = typename T::reference;

	MultiContainer()
	{}

	template<typename ... Args>
	MultiContainer(Args ... args)
	{
		(add(args),...);
	}
	
	template<typename CI, typename TI, typename U>
	static auto begin(U &self)
	{
		auto &c = self.c;
		
		CI end = c.end();
		CI i = c.begin();
		TI j;
		
		auto ended = c.empty();
		if (!ended)
		{
			j = (*i)->begin();
		}
		
		return MultiIterator<T,CI,TI>(
			std::move(end),
			std::move(i),
			std::move(j),
			ended
		);
	}
	
	template<typename CI, typename TI, typename U>
	static auto end(U &self)
	{
		auto &c = self.c;
		
		CI end = c.end();
		CI i = c.end();
		TI j;
		
		return MultiIterator<T,CI,TI>(
			std::move(end),
			std::move(i),
			std::move(j),
			true
		);
	}
	
	auto begin()
	{
		return begin<typename C::iterator, typename T::iterator>(*this);
	}
	
	auto end()
	{
		return end<typename C::iterator, typename T::iterator>(*this);
	}
	
	auto begin() const
	{
		return begin<typename C::const_iterator, typename T::const_iterator>(*this);
	}
	
	auto end() const
	{
		return end<typename C::const_iterator, typename T::const_iterator>(*this);
	}

	void add(T *t)
	{
		c.push_back(t);
	}
	
	void remove(T *t)
	{
		auto i = std::find(c.begin(), c.end(), t);
		if (i != c.end())
		{
			std::swap(*i, c.back());
			c.pop_back();
		}
	}
	
	auto container_size() const
	{
		return c.size();
	}
	
	auto size() const
	{
		size_t size_ = 0;
		for (auto *v : c)
			size_ += v->size();
			
		return size_;
	}
	
	auto empty() const
	{
		return size() == 0;
	}
	
	auto &at(size_t i) const
	{
		for (auto *v : c)
		{
			size_t size_ = v->size();
			if (i < size_)
				return v->at(i);
				
			i -= size_;
		}
			
		debug_assert(false);
		throw exceptions::OutOfBounds();
	}
	
	auto &operator[](size_t i)
	{
		return at(i);
	}
} ;

} // namespace