#pragma once

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

#include <jvalidate/detail/anchor.h>
#include <jvalidate/detail/on_block_exit.h>
#include <jvalidate/detail/reference.h>
#include <jvalidate/uri.h>

namespace jvalidate::detail {
/**
 * @brief Starting with Draft 2019-09, it is possible to set $recursiveAnchor
 * and $recursiveRef options in schemas. In Draft 2020-12 this changes to the
 * more powerful/generic $dynamicAnchor and $dynamicRef.
 *
 * The rules of handling these anchors is that we maintain a stack of all of the
 * loaded anchors for each given name in the order that they are loaded. But the
 * resolved reference that they point to is the first anchor by that name to be
 * loaded. This means that we can create recursive/self-referential chains.
 *
 * When we encounter the appropriate $dynamicRef/$recursiveRef tag, we fetch
 * the most appropriate anchored location - which is usually the prior mentioned
 * first path that registered the anchor.
 */
class DynamicReferenceContext {
private:
  std::deque<URI> sources_;
  std::map<Anchor, std::deque<std::optional<Reference>>> data_;

public:
  /**
   * @brief Add all dynamic anchors contained in a given document (as defined
   * by a common URI) to the current stack, pointing them to the earliest loaded
   * parent reference and unregistering all of the anchors that are in context,
   * but not in this specific document.
   */
  OnBlockExit scope(URI const & source, std::map<Anchor, Reference> const & frame) {
    // No-Op loading, for convenience
    if (frame.empty() && data_.empty()) {
      return nullptr;
    }

    sources_.push_back(source);
    for (auto const & [k, v] : frame) {
      // If we have not currently registered this anchor, use the input
      // reference path, else use the first reference path registered in the
      // stack.
      data_[k].push_back(data_[k].empty() ? v : data_[k].front());
    }

    // For all of the anchors that are not being pushed onto the stack, push a
    // nullopt onto the stack (as well as ensuring that all stacks are
    // equal-sized). This allows us to disable certain anchors in a given
    // document (i.e. that doc does not define a specific $dynamicAnchor).
    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);
      }
    }

    // Scope object the pops all of the elements on this object, due to the
    // equal-size promise of the above while loop, we can just blindly loop
    // through all elements to pop instead of dealing with
    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;
        }
      }
    };
  }

  /**
   * @brief Is the given anchor in the current $dynamicRef lookup context
   * (including suppressed anchors). This check is necessary in part because we
   * permit using $dynamicRef to refer regular $anchor objects if there is no
   * $dynamicAnchor in the current context.
   *
   * TODO(samjaffe): May be able to add a nullopt check...
   */
  bool contains(Anchor const & key) const { return data_.contains(key); }

  /**
   * @brief Safely fetch the closest matching $dynamicAnchor to the given
   * arguments. Because $dynamicRef permits including a URI, it is techinically
   * possible to "jump" to an anchor that is not the top-level one, this can
   * be useful if the current scope does not generate a bookmark $dynamicAnchor.
   *
   * @param source The owning source, which is either the URI of the currently
   * operating schema document, or a URI specified in the $dynamicRef value.
   * Using this information lets us jump past suppressed anchors by explicitly
   * stating the owning context.
   *
   * @param key The actual anchor being searched for.
   */
  std::optional<Reference> lookup(URI const & source, Anchor const & key) const {
    if (auto it = data_.find(key); it != data_.end()) {
      return it->second.at(index_of(source));
    }
    return std::nullopt;
  }

  /**
   * @brief Finds the (index of the) dynamic anchor directly associated with the
   * given URI; or the final registered anchor.
   */
  size_t index_of(URI const & source) const {
    // Start at the end because most commonly source will refer to the currently
    // operating schema, which is also going to be the top item in the sources
    // stack.
    for (size_t i = sources_.size(); i-- > 0;) {
      if (sources_[i] == source) {
        return i;
      }
    }
    return sources_.size() - 1;
  }

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