#pragma once

#include <map>
#include <unordered_set>

#include <jvalidate/compat/enumerate.h>
#include <jvalidate/detail/anchor.h>
#include <jvalidate/detail/dynamic_reference_context.h>
#include <jvalidate/detail/expect.h>
#include <jvalidate/detail/on_block_exit.h>
#include <jvalidate/detail/out.h>
#include <jvalidate/detail/parser_context.h>
#include <jvalidate/detail/pointer.h>
#include <jvalidate/detail/reference.h>
#include <jvalidate/detail/reference_cache.h>
#include <jvalidate/detail/vocabulary.h>
#include <jvalidate/document_cache.h>
#include <jvalidate/enum.h>
#include <jvalidate/forward.h>
#include <jvalidate/uri.h>

namespace jvalidate::detail {
/**
 * @brief An object responsible for owning/managing the various documents,
 * references, and related functionality for ensuring that we properly construct
 * things.
 *
 * In order to support this we store information on:
 * - A {@see jvalidate::detail::ReferenceCache} that maps various absolute
 *   Reference paths to their Canonical forms.
 * - "Vocabularies", which describe the the set of legal keywords for
 *   constraint parsing.
 * - "Anchor Locations", a non-owning store of the Adapters associated with
 *   "$id"/"$anchor" tags to allow quick lookups without having to re-walk the
 *   document.
 * - "Dynamic Anchors", a list of all of the "$dynamicAnchor" tags that exist
 *   under a given "$id" tag, and those bindings which are active in the current
 *   scope.
 *
 * @tparam A The adapter type being operated upon
 */
template <Adapter A> class ReferenceManager {
private:
  static inline std::map<std::string_view, schema::Version> const g_schema_ids{
      {"json-schema.org/draft-03/schema", schema::Version::Draft03},
      {"json-schema.org/draft-04/schema", schema::Version::Draft04},
      {"json-schema.org/draft-06/schema", schema::Version::Draft06},
      {"json-schema.org/draft-07/schema", schema::Version::Draft07},
      {"json-schema.org/draft/2019-09/schema", schema::Version::Draft2019_09},
      {"json-schema.org/draft/2020-12/schema", schema::Version::Draft2020_12},
  };

  ConstraintFactory<A> const & constraints_;
  DocumentCache<A> & external_;

  ReferenceCache references_;

  std::map<schema::Version, Vocabulary<A>> vocabularies_;
  std::map<URI, Vocabulary<A>> user_vocabularies_;

  std::map<RootReference, A> roots_;

  std::map<URI, std::map<Anchor, Reference>> dynamic_anchors_;
  DynamicReferenceContext active_dynamic_anchors_;

public:
  /**
   * @brief Construct a new ReferenceManager around a given root schema
   *
   * @param external A cache/loader of external documents. Due to the way that
   * {@see jvalidate::Schema} is implemented, the cache may have the same
   * lifetime as this object, despite being owned by mutable reference.
   *
   * @param root The root schema being operated on.
   *
   * @param version The version of the schema being used for determining the
   * base vocabulary to work with (see the definition of schema::Version for
   * more details on how the base vocabulary changes).
   *
   * @param constraints A factory for turning JSON schema information into
   * constraints.
   */
  ReferenceManager(DocumentCache<A> & external, A const & root, schema::Version version,
                   ConstraintFactory<A> const & constraints)
      : external_(external), constraints_(constraints), roots_{{{}, root}} {
    prime(root, {}, &vocab(version));
  }

  /**
   * @brief Turn a schema version into a vocabulary, ignoring user-defined
   * vocabularies
   *
   * @param version The schema version
   *
   * @returns The default vocabulary for a given draft version
   */
  Vocabulary<A> const & vocab(schema::Version version) {
    if (not vocabularies_.contains(version)) {
      vocabularies_.emplace(version, constraints_.keywords(version));
    }
    return vocabularies_.at(version);
  }

  /**
   * @brief Fetch the vocabulary information associated with a given "$schema"
   * tag. Unlike the enum version of this function, we can also load
   * user-defined schemas using the ReferenceCache object, if supported. This
   * allows us to define custom constraints or remove some that we want to
   * forbid.
   *
   * @param schema The location of the schema being fetched
   *
   * @returns If schema is a draft version - then one of the default
   * vocabularies, else a user-schema is loaded.
   */
  Vocabulary<A> const & vocab(URI schema) {
    if (auto it = g_schema_ids.find(schema.resource()); it != g_schema_ids.end()) {
      return vocab(it->second);
    }

    if (auto it = user_vocabularies_.find(schema); it != user_vocabularies_.end()) {
      return it->second;
    }

    std::optional<A> external = external_.try_load(schema);
    EXPECT_M(external.has_value(), "Unable to load external meta-schema " << schema);
    EXPECT_M(external->type() == adapter::Type::Object, "meta-schema must be an object");

    auto metaschema = external->as_object();
    // All user-defined schemas MUST have a parent schema they point to
    // Furthermore - in order to be well-formed, the schema chain must
    // eventually point to one of the draft schemas. However - if a metaschema
    // ends up in a recusive situation (e.g. A -> B -> A), it will not fail in
    // the parsing step, but instead produce a malformed Schema object for
    // validation.
    EXPECT_M(metaschema.contains("$schema"),
             "user-defined meta-schema must reference a base schema");

    // Initialize first to prevent recursion
    Vocabulary<A> & parent = user_vocabularies_[schema];
    parent = vocab(URI(metaschema["$schema"].as_string()));

    if (metaschema.contains("$vocabulary")) {
      // This is a silly thing we have to do because rather than have some kind
      // of annotation/assertion divide marker for the format constraint, we
      // instead use true/false in Draft2019-09, and have format-assertion/
      // format-annotation vocabularies in Draft2020-12.
      auto [keywords, vocabularies] = extract_keywords(metaschema["$vocabulary"].as_object());
      parent.restrict(keywords, vocabularies);
    }

    return parent;
  }

  /**
   * @brief Load the current location into the stack of dynamic ref/anchors so
   * that we are able to properly resolve them (e.g. because an anchor got
   * disabled).
   *
   * @param ref The current parsing location in the schema, which should
   * correspond with an "$id" tag.
   *
   * @returns A scope object that will remove this set of dynamic ref/anchor
   * resolutions from the stack when it exits scope.
   */
  auto dynamic_scope(Reference const & ref) {
    URI const uri =
        ref.pointer().empty() ? ref.uri() : references_.relative_to_nearest_anchor(ref).uri();
    return active_dynamic_anchors_.scope(uri, dynamic_anchors_[uri]);
  }

  /**
   * @breif "Load" a requested document reference, which may exist in the
   * current document, or in an external one.
   *
   * @param ref The location to load. Since there is no guarantee of direct
   * relation between the current scope and this reference, we treat this like a
   * jump.
   *
   * @param vocab The current vocabulary being used for parsing. It may be
   * changed when loading the new reference if there is a "$schema" tag at the
   * root.
   *
   * @returns The schema corresponding to the reference, if it can be located.
   * As long as ref contains a valid URI/Anchor, we will return an Adapter, even
   * if that adapter might point to a null JSON.
   */
  std::optional<A> load(Reference const & ref, Vocabulary<A> const * vocab) {
    if (auto it = roots_.find(ref.root()); it != roots_.end()) {
      return ref.pointer().walk(it->second);
    }

    std::optional<A> external = external_.try_load(ref.uri());
    if (not external) {
      return std::nullopt;
    }

    references_.emplace(ref.uri());
    prime(*external, ref, vocab);

    // May have a sub-id that we map to
    if (auto it = roots_.find(ref.root()); it != roots_.end()) {
      return ref.pointer().walk(it->second);
    }

    // Will get called if the external schema does not declare a root id?
    return ref.pointer().walk(*external);
  }

  /**
   * @brief Transform a reference into its "canonical" form, in the context of
   * the calling context (parent).
   *
   * @param ref The value of a "$ref" or "$dynamicRef" token, that is being
   * looked up.
   *
   * @param parent The current lexical scope being operated in.
   *
   * @param dynamic_reference As an input, indicates that we are requesting a
   * dynamic reference instead of a normal $ref.
   * As an output, indicates that we effectively did resolve a dynamicRef and
   * therefore should alter the dynamic scope in order to prevent infinite
   * recursions in schema parsing.
   *
   * @returns ref, but in its canonical/lexical form.
   */
  Reference canonicalize(Reference const & ref, Reference const & parent,
                         inout<bool> dynamic_reference) {
    URI const uri = [this, &ref, &parent]() {
      // If there are no URIs involed (root schema does not set "$id")
      // then we don't need to do anything clever
      if (ref.uri().empty() && parent.uri().empty()) {
        return references_.actual_parent_uri(parent);
      }

      // At least one of ref and parent have a real URI/"$id" value. If it has a
      // "root" (e.g. file:// or http://), then we don't need to do any clever
      // alterations to identify the root.
      URI uri = ref.uri().empty() ? parent.uri() : ref.uri();
      if (not uri.is_rootless()) {
        return uri;
      }

      // Now we need to compute that URI into the context of its parent, such
      // as if ref := "file.json" and
      // parent := "http://localhost:8000/schemas/root.json"
      URI base = references_.actual_parent_uri(parent);
      EXPECT_M(base.resource().rfind('/') != std::string::npos,
               "Unable to deduce root for relative uri " << uri << " (" << base << ")");
      if (not uri.is_relative()) {
        return base.root() / uri;
      }
      if (auto br = base.resource(), ur = uri.resource();
          br.ends_with(ur) && br[br.size() - ur.size() - 1] == '/') {
        return base;
      }
      return base.parent() / uri;
    }();

    // This seems unintuitive, but we generally want to avoid providing a URI
    // when looking up dynamic references, unless they are explicitly asked for.
    URI const dyn_uri = ref.uri().empty() ? URI() : uri;
    if (std::optional dynref = dynamic(dyn_uri, ref, dynamic_reference)) {
      return *dynref;
    }

    dynamic_reference = dynamic_reference || active_dynamic_anchors_.empty();

    // Relative URI, not in the HEREDOC (or we set an $id)
    if (ref.uri().empty() and ref.anchor().empty()) {
      return Reference(references_.relative_to_nearest_anchor(parent).root(), ref.pointer());
    }

    return Reference(uri, ref.anchor(), ref.pointer());
  }

private:
  /**
   * @brief Locate the dynamic reference being requested (if it is being
   * requested).
   *
   * @param uri The dynamic reference uri being requested, generally empty.
   *
   * @param ref The value of a "$ref" or "$dynamicRef" token, that is being
   * looked up. Primarily used for the anchor value, which is relevant for
   * $dynamicRef/$dynamicAnchor.
   *
   * @param dynamic_reference As an input, indicates that we are requesting a
   * dynamic reference instead of a normal $ref.
   * As an output, indicates that we effectively did resolve a dynamicRef and
   * therefore should alter the dynamic scope in order to prevent infinite
   * recursions in schema parsing.
   *
   * @returns If there is a dynamic reference for the requested anchor, we
   * return it.
   */
  std::optional<Reference> dynamic(URI const & uri, Reference const & ref,
                                   inout<bool> dynamic_reference) {
    bool const anchor_is_dynamic = active_dynamic_anchors_.contains(ref.anchor());
    if (not dynamic_reference) {
      // A normal $ref to an $anchor that matches a $dynamicAnchor breaks the
      // dynamic recursion pattern. This requires that we are not looking for a
      // subschema of the anchor AND that we are not targetting an anchor in a
      // different root document.
      dynamic_reference = (anchor_is_dynamic && ref.uri().empty() && ref.pointer().empty());
      return std::nullopt;
    }

    OnBlockExit scope;
    if (not ref.uri().empty() && anchor_is_dynamic) {
      // Register the scope of this (potential) $dynamicAnchor BEFORE we attempt
      // to enter the reference, in case we end up pointing to an otherwise
      // suppressed $dynamicAnchor in a higher scope.
      scope = dynamic_scope(Reference(uri));
    }

    return active_dynamic_anchors_.lookup(uri, ref.anchor());
  }

  /**
   * @brief Prepare a newly loaded document, importing schema information,
   * ids, anchors, and dynamic anchors recursively.
   *
   * @param json The document being loaded
   *
   * @param vocab The vocabulary of legitimate keywords to iterate through to
   * locate ids etc.
   */
  void prime(Adapter auto const & json, Reference where, Vocabulary<A> const * vocab) {
    if (json.type() != adapter::Type::Object) {
      return;
    }

    auto schema = json.as_object();
    // Update vocabulary to the latest form
    if (schema.contains("$schema")) {
      vocab = &this->vocab(URI(schema["$schema"].as_string()));
    }

    // Load ids, anchors, etc.
    prime_roots(where, vocab->version(), json);

    // Recurse through the document
    for (auto const & [key, value] : schema) {
      if (not vocab->is_keyword(key)) {
        continue;
      }
      switch (value.type()) {
      case adapter::Type::Array: {
        // Recurse through array-type schemas, such as anyOf, allOf, and oneOf
        // we don't actually check that the key is one of those, because if we
        // do something stupid like "not": [] then the parsing phase will return
        // an error.
        for (auto const & [index, elem] : detail::enumerate(value.as_array())) {
          prime(elem, where / key / index, vocab);
        }
        break;
      }
      case adapter::Type::Object:
        // Normal schema-type data such as not, additionalItems, etc. hold a
        // schema as their immidiate child.
        if (not vocab->is_property_keyword(key)) {
          prime(value, where / key, vocab);
          break;
        }
        // Special schemas are key-value stores, where the key is arbitrary and
        // the value is the schema. Therefore we need to skip over the props.
        for (auto const & [prop, elem] : value.as_object()) {
          prime(elem, where / key / prop, vocab);
        }
      default:
        break;
      }
    }
  }

  /**
   * @brief Optionally register any root document at this location, as
   * designated by things like the "$id" and "$anchor" tags.
   *
   * @param where The current lexical location in the schema - if there is an
   * id/anchor tag, then we overwrite this value with the newly indicated root.
   *
   * @param version The current schema version - used to denote the name of the
   * id tag, whether anchors are available, and how dynamic anchors function
   * (Draft2019-09's recursiveAnchor vs. Draft2020-12's dynamicAnchor).
   *
   * @param json The document being primed.
   */
  void prime_roots(Reference & where, schema::Version version, A const & json) {
    std::string const id = version <= schema::Version::Draft04 ? "id" : "$id";
    auto const schema = json.as_object();

    RootReference root = where.root();
    if (schema.contains(id)) {
      root = RootReference(schema[id].as_string());
      if (root.uri().empty()) {
        root = RootReference(where.uri(), root.anchor());
      } else if (not root.uri().is_rootless() || where.uri().empty()) {
        // By definition - rooted URIs cannot be relative
      } else if (root.uri().is_relative()) {
        root = RootReference(where.uri().parent() / root.uri(), root.anchor());
      } else {
        root = RootReference(where.uri().root() / root.uri(), root.anchor());
      }

      roots_.emplace(root, json);
      where = references_.emplace(where, root);
    }

    // $anchor and its related keywords were introduced in Draft 2019-09
    if (version < schema::Version::Draft2019_09) {
      return;
    }

    if (schema.contains("$anchor")) {
      root = RootReference(root.uri(), Anchor(schema["$anchor"].as_string()));
      roots_.emplace(root, json);
      where = references_.emplace(where, root);
    }

    // Unfortunately - $recursiveAnchor and $dynamicAnchor use very different
    // handling mechanisms, so it is not convenient to merge together
    if (version == schema::Version::Draft2019_09 && schema.contains("$recursiveAnchor") &&
        schema["$recursiveAnchor"].as_boolean()) {
      Anchor anchor;
      root = RootReference(root.uri(), anchor);

      roots_.emplace(root, json);
      where = references_.emplace(where, root);

      if (Reference & dynamic = dynamic_anchors_[root.uri()][anchor];
          dynamic == Reference() || where < dynamic) {
        dynamic = where;
      }
    }

    if (schema.contains("$dynamicAnchor") && version > schema::Version::Draft2019_09) {
      Anchor anchor(schema["$dynamicAnchor"].as_string());
      root = RootReference(root.uri(), anchor);

      roots_.emplace(root, json);
      where = references_.emplace(where, root);

      if (Reference & dynamic = dynamic_anchors_[root.uri()][anchor];
          dynamic == Reference() || where < dynamic) {
        dynamic = where;
      }
    }
  }

  /**
   * @brief Extract the supported keywords of a given selection of vocabularies
   *
   * @param vocabularies A map of the form (VocabularyURI => Enabled)
   *
   * @returns A pair containing:
   * - All of the enabled keywords in the vocabulary
   * - The list of enabled vocabulary metaschema (used for is_format_assertion)
   */
  auto extract_keywords(ObjectAdapter auto const & vocabularies) const
      -> std::pair<std::unordered_map<std::string, bool>, std::unordered_set<std::string>> {
    std::unordered_map<std::string, bool> keywords;
    std::unordered_set<std::string> vocab_docs;
    for (auto [vocab, required] : vocabularies) {
      size_t n = vocab.find("/vocab/");
      vocab_docs.emplace(vocab.substr(n));
      vocab.replace(n, 7, "/meta/");

      auto vocab_object = external_.try_load(URI(vocab));
      auto it = vocab_object->as_object().find("properties");
      if (it == vocab_object->as_object().end()) {
        continue;
      }

      for (auto const & [keyword, _] : it->second.as_object()) {
        keywords.emplace(keyword, required.as_boolean());
      }
    }
    return std::make_pair(keywords, vocab_docs);
  }
};
}