#pragma once
#include <stack>
#include <iterator>

#include "iterators/BaseIterator.h"

template <typename T>
class DFSIterator : public BaseIterator<T> {
public:
	explicit DFSIterator(T* ptr) : BaseIterator<T>(ptr) {
		if (!this->current())
			return;

		for (auto it = this->current()->children().rbegin();
		     it != this->current()->children().rend(); ++it)
			s.push({it->get(), 1});
	}

	IIterator<T>& operator++() override {
		advance();
		return *this;
	}

protected:
	std::stack<std::pair<T*, size_t>> s;

	void advance() {
		if (s.empty()) {
			this->set_current(nullptr);
			return;
		}

		auto [node, lvl] = s.top();
		s.pop();

		this->set_current(node);
		this->set_level(lvl);

		for (auto it = this->current()->children().rbegin();
		     it != this->current()->children().rend(); ++it)
			s.push({it->get(), lvl + 1});
	}
};