#pragma once

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

#include <jvalidate/detail/on_block_exit.h>

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

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

    sources_.push_back(source);
    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);
      }
      while (stack.size() < sources_.size()) {
        stack.push_front(std::nullopt);
      }
    }

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

  bool contains(Key const & key) const { return data_.contains(key); }

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

  size_t index_of(Source const & source) const {
    for (size_t i = sources_.size(); i-- > 0;) {
      if (sources_[i] == source) {
        return i;
      }
    }
    return sources_.size() - 1;
  }

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