#pragma once

#include <cassert>
#include <vector>
#include <memory>

namespace td {

template <typename TDerived, typename TReturn, typename... TArgs>
class StateStack {
  public:
	class State {
	  public:
		State(TDerived& a_StateStack) : m_StateStack(a_StateStack) {}
		virtual ~State() {}

		virtual TReturn Update(TArgs... args) = 0;

	  protected:
		/**
		 * \brief Called when enabled (can be used to connect signals)
		 */
		virtual void OnEnable() {}

		/**
		 * \brief Called when disabled (can be used to disconnect signals)
		 */
		virtual void OnDisable() {}

	  protected:
		TDerived& m_StateStack;

		friend class StateStack;
	};

	StateStack() {}
	StateStack(StateStack&&) = default;
	virtual ~StateStack() {}

	virtual TReturn Update(TArgs... args) {
		assert(!m_States.empty() && "You must push at least one state before updating !");
		return m_States.back()->Update(args...);
	}

	void PopState() {
		assert(!m_States.empty() && "You must push at least one state before poping !");
		m_States.back()->OnDisable();
		m_States.pop_back();
		if (m_States.empty())
			return;
		m_States.back()->OnEnable();
	}

	template <typename T, typename... Args>
	T* PushState(Args&&... args) {
		if (!m_States.empty())
			m_States.back()->OnDisable();
		auto newState = std::make_unique<T>(static_cast<TDerived&>(*this), std::forward<Args>(args)...);
		newState->OnEnable();
		m_States.push_back(std::move(newState));
		return static_cast<T*>(newState.get());
	}

  private:
	std::vector<std::unique_ptr<State>> m_States;
};


}  // namespace td