#pragma once

#include <deque>
#include <map>
#include <optional>

#include <jvalidate/detail/on_block_exit.h>

namespace jvalidate::detail {
template <typename Key, typename Value> class ContextStack {
private:
  std::map<Key, std::deque<std::optional<Value>>> data_;

public:
  OnBlockExit scope(std::map<Key, Value> const & frame) {
    if (frame.empty() && data_.empty()) {
      return nullptr;
    }

    for (auto const & [k, v] : frame) {
      data_[k].push_back(data_[k].empty() ? v : data_[k].front());
    }

    for (auto & [k, stack] : data_) {
      if (not frame.contains(k)) {
        stack.push_back(std::nullopt);
      }
    }

    return [this]() {
      for (auto it = data_.begin(); it != data_.end();) {
        if (it->second.size() == 1) {
          it = data_.erase(it);
        } else {
          it->second.pop_back();
        }
      }
    };
  }

  std::optional<Value> lookup(Key const & key) const {
    if (auto it = data_.find(key); it != data_.end()) {
      return it->second.back();
    }
    return std::nullopt;
  }

  bool empty() const { return data_.empty(); }
};
}