#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), level_(0) {
		if (!this->current)
			return;

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

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

	virtual size_t level() const {
		return level_;
	}

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

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

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

		this->current = node;
		level_ = lvl;

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