#pragma once

#include <functional>
#include <map>
#include <optional>

#include <jvalidate/detail/pointer.h>
#include <jvalidate/detail/reference.h>

namespace jvalidate::detail {
/**
 * @brief An bidirectional cache of absolute references
 * (URI + Anchor + JSON-Pointer) to root references (URI + Anchor). An object of
 * this sort is necessary for a couple of reasons:
 *
 * 1. It is possible to have more than one absolute reference map to the same
 *    root reference.
 * 2. We need to employ some special handling to properly identify the nearest
 *    RootReference that owns a given absolute Reference - given that that ref
 *    may not actually be anchored to the RootReference at the time.
 */
class ReferenceCache {
private:
  std::map<Reference, RootReference, std::greater<>> to_anchor_;
  std::map<RootReference, Reference, std::greater<>> to_absolute_;

public:
  /**
   * @brief Register the entry-point of a given schema-document. This function
   * should be called exactly one time for each "$id" tag in a given schema that
   * is being loaded.
   *
   * In principle, we should also perform a uniqueness check when calling
   * to_absolute_.emplace.
   */
  void emplace(URI const & root) {
    to_absolute_.emplace(root, root);
    to_anchor_.emplace(root, root);
  }

  /**
   * @brief Link together the absolute and root references of a document, as
   * well as recursively walk through all possible parent URIs that are already
   * stored in this cache.
   *
   * Therefore, this function will add "exactly 1" mapping to the to_absolute_
   * map, and "at least 1, but no more than to_absolute_.size()" mappings to
   * to_anchor_, representing all of the parent reference paths that link to the
   * newly added RootReference
   *
   * @param absolute An absolute JSON Reference, which either contains no
   * RootReference component, or contains the previous traversed RootReference,
   * as defined by the $id, $anchor, $recursiveAnchor, or $dynamicAnchor tags.
   *
   * @param canonical The current RootReference being operated on from the
   * tags listed above.
   *
   * For example, if we have a json document like:
   * @code{.json}
   * {
   *   "$id": "A",
   *   "$defs": {
   *     "Example": {
   *       "$id": "B"
   *     }
   *   }
   * }
   * @endcode
   *
   * then we would end up calling this function with the arguments:
   *    absolute="#",                canonical="A#"
   *    absolute="A#/$defs/Example", canonical="B#"
   */
  Reference emplace(Reference const & absolute, RootReference const & canonical) {
    for (Reference where = absolute; not where.pointer().empty();) {
      // Recursively add all possible alternative paths that are equivalent to
      // this absolute path as valid mappings to this canonical one.
      auto it = to_absolute_.lower_bound(where.root());
      if (it == to_absolute_.end() || where.root() == it->second.root()) {
        break;
      }

      if (auto ptr = it->second.pointer(); where.pointer().starts_with(ptr)) {
        where = it->second / where.pointer().relative_to(ptr);
        to_anchor_.emplace(where, canonical);
      } else {
        break;
      }
    }

    to_anchor_.emplace(absolute, canonical);
    to_absolute_.emplace(canonical, absolute);

    return Reference(canonical);
  }

  /**
   * @brief Identifies the nearest RootReference that is associated with the
   * input.
   *
   * @param ref An arbitrary reference that we want to locate the nearest root
   * for.
   *
   * @param for_parent_reference A flag indicating if we should prohibit exact
   * matches of reference. For example, suppose that we have the same bindings
   * as in the above method comment:
   *    absolute="#",                canonical="A#"
   *    absolute="A#/$defs/Example", canonical="B#"
   * If I request `relative_to_nearest_anchor("A#/$defs/Example", false)` then
   * it will return `B#` as the associated RootReference, because we have stored
   * that exact mapping in our absolute path to anchor cache.
   *
   * On the other hand - suppose that we want to ensure that we've acquired
   * strict parent of the current reference.
   * `relative_to_nearest_anchor("A#/$defs/Example", true)` would say "an anchor
   * cannot be its own parent, therefore we cannot resolve to B#".
   *
   * @returns ref, recalculated to be relative_to its nearest parent root, if
   * one is available.
   */
  Reference relative_to_nearest_anchor(Reference const & ref,
                                       bool for_parent_reference = false) const {
    auto it = for_parent_reference ? to_anchor_.upper_bound(ref) : to_anchor_.lower_bound(ref);
    if (it == to_anchor_.end()) {
      return ref;
    }

    auto const & [absolute, anchor] = *it;
    if (not ref.pointer().starts_with(absolute.pointer())) {
      // We've undershot our reference and landed at a cousin/neighbor node
      return ref;
    }

    return Reference(anchor, ref.pointer().relative_to(absolute.pointer()));
  }

  /**
   * @brief Deduces the URI part of the actual parent of this node, utilizing
   * the "an anchor cannot be its own parent" rule described above.
   *
   * @param parent An arbitrarily constructed reference to the parent context
   * of some other reference we are operating on.
   *
   * @returns The URI of the nearest non-equal parent if it exists and/or there
   * is a URI part in parent. If there is no URI part, we check if there is a
   * URI for the root (input) schema.
   */
  URI actual_parent_uri(detail::Reference const & parent) const {
    // TODO(samjaffe): Think about this some more - there's something awkward here
    URI uri = relative_to_nearest_anchor(parent, true).uri();

    if (not uri.empty() || not parent.uri().empty()) {
      return uri;
    }

    // This is a special case because we prohibit exact matches in the above
    // relative_to_nearest_anchor call. Since we can only reach this line if
    // BOTH uri and parent.uri() are empty - that means that the appropriate
    // parent is the root document, which might have marked its URI with an
    // $id tag.
    if (auto it = to_anchor_.find(Reference()); it != to_anchor_.end()) {
      return it->second.uri();
    }

    return uri;
  }
};
}