#pragma once

#include <memory>
#include <type_traits>
#include <unordered_map>
#include <vector>

#include <jvalidate/adapter.h>
#include <jvalidate/constraint.h>
#include <jvalidate/detail/anchor.h>
#include <jvalidate/detail/expect.h>
#include <jvalidate/detail/on_block_exit.h>
#include <jvalidate/detail/parser_context.h>
#include <jvalidate/detail/pointer.h>
#include <jvalidate/detail/reference.h>
#include <jvalidate/detail/reference_manager.h>
#include <jvalidate/document_cache.h>
#include <jvalidate/enum.h>
#include <jvalidate/forward.h>

namespace jvalidate::schema {
/**
 * @brief The real "Schema" class, representing a resolved node in a schema
 * object. Each node is analogous to one layer of the schema json, and can
 * represent either a "rejects all" schema, an "accepts all" schema, or a
 * schema that has some selection of constraints and other features.
 */
class Node {
private:
  // Annotations for this schema...
  std::string description_;

  // The default value to apply to an object if if does not exist - is invoked
  // by the parent schema node, rather than this node itself.
  std::unique_ptr<adapter::Const const> default_{nullptr};

  // Rejects-all can provide a custom reason under some circumstances.
  std::optional<std::string> rejects_all_;

  // Actual constraint information
  std::optional<schema::Node const *> reference_{};
  std::unordered_map<std::string, std::unique_ptr<constraint::Constraint>> constraints_{};
  std::unordered_map<std::string, std::unique_ptr<constraint::Constraint>> post_constraints_{};

public:
  Node() = default;
  /**
   * @brief Construct a schema that rejects all values, with a custom reason
   *
   * @param A user-safe justification of why this schema rejects everything.
   * Depending on the compiler settings, this might be used to indicate things
   * such as attempting to load a non-existant schema.
   */
  Node(std::string const & rejection_reason) : rejects_all_(rejection_reason) {}

  /**
   * @brief Actually initialize this schema node. Unfortunately, we cannot use
   * RAII for initializing this object because of certain optimizations and
   * guardrails make reference captures breakable.
   *
   * @param context The currently operating context, including the actual JSON
   * document being parsed at this moment.
   */
  template <Adapter A> void construct(detail::ParserContext<A> context);

  bool is_pure_reference() const {
    return reference_ && constraints_.empty() && post_constraints_.empty() && not default_;
  }
  bool accepts_all() const {
    return not reference_ && constraints_.empty() && post_constraints_.empty();
  }
  std::optional<std::string> const & rejects_all() const { return rejects_all_; }
  std::optional<schema::Node const *> reference_schema() const { return reference_; }

  bool requires_result_context() const { return not post_constraints_.empty(); }
  auto const & constraints() const { return constraints_; }
  auto const & post_constraints() const { return post_constraints_; }

  adapter::Const const * default_value() const { return default_.get(); }

private:
  /**
   * @brief Resolve any dynamic anchors that are children of the current schema
   * (if this is the root node of a schema). If it is not a root node (does not
   * define "$id"), then this function does nothing.
   *
   * @tparam A The Adapter type for the JSON being worked with.
   *
   * @param context The currently operating context, including the actual JSON
   * document being parsed at this moment.
   *
   * @returns If this is a root schema - a scope object to pop the dynamic scope
   */
  template <Adapter A> detail::OnBlockExit resolve_anchor(detail::ParserContext<A> const & context);

  /**
   * @brief Resolves/embeds referenced schema information into this schema node.
   *
   * @tparam A The Adapter type for the JSON being worked with.
   *
   * @param context The currently operating context, including the actual JSON
   * document being parsed at this moment.
   *
   * @returns true iff there was a reference tag to follow
   */
  template <Adapter A> bool resolve_reference(detail::ParserContext<A> const & context);
};
}

namespace jvalidate {

class Schema : public schema::Node {
private:
  friend class schema::Node;
  template <Adapter A> friend class detail::ParserContext;

private:
  schema::Node accept_;
  schema::Node reject_{"always false"};

  // An owning cache of all created schemas. Avoids storing duplicates such as
  // the "always-true" schema, "always-false" schema, and schemas whose only
  // meaningful field is "$ref", "$recursiveRef", or "$dynamicRef".
  std::map<detail::Reference, schema::Node> cache_;

  // A non-owning cache of all schemas, including duplcates where multiple
  // References map to the same underlying schema.
  std::map<detail::Reference, schema::Node const *> alias_cache_;

public:
  /**
   * @brief Construct a new schema. All other constructors of this type may be
   * considered syntactic sugar for this constructor.
   *
   * As such, the true signature of this class's contructor is:
   *
   * Schema(Adapter| JSON
   *        [, schema::Version]
   *        [, URIResolver<A> | DocumentCache<A> &]
   *        [, ConstraintFactory<A> const &])
   *
   * as long as the order of arguments is preserved - the constructor will work
   * no matter which arguments are ignored. The only required argument being
   * the JSON object/Adapter.
   *
   * @param json An adapter to a json object
   *
   * @param version The json-schema draft version that all schemas will prefer
   *
   * @param external An object capable of resolving URIs, and turning them into
   * Adapter objects. Holds a cache and so must be mutable.
   *
   * @param factory An object that manuafactures constraints - allows the user
   * to provide custom extensions or even modify the behavior of existing
   * keywords by overridding the virtual accessor function(s).
   */
  template <Adapter A>
  Schema(A const & json, schema::Version version, DocumentCache<A> & external,
         ConstraintFactory<A> const & factory = {}) {
    // Prevent unintialized data caches
    if (version >= schema::Version::Draft06 && json.type() == adapter::Type::Boolean) {
      schema::Node::operator=(std::move(json.as_boolean() ? accept_ : reject_));
      return;
    }

    detail::ReferenceManager<A> ref(external, json, version, factory);
    detail::ParserContext<A> root{*this, json, &ref.vocab(version), ref};

    root.where = root.dynamic_where = ref.canonicalize({}, {}, false);
    construct(root);
  }

  /**
   * @param json An adapter to a json schema
   *
   * @param version The json-schema draft version that all schemas will prefer
   *
   * @param external An object capable of resolving URIs, and turning them into
   * Adapter objects. Holds a cache and so must be mutable. If this constructor
   * is called, then it means that the cache is a one-off object, and will not
   * be reused.
   */
  template <Adapter A, typename... Args>
  Schema(A const & json, schema::Version version, DocumentCache<A> && external, Args &&... args)
      : Schema(json, version, external, std::forward<Args>(args)...) {}

  /**
   * @param json An adapter to a json schema
   *
   * @param version The json-schema draft version that all schemas will prefer
   *
   * @param resolve A function capable of resolving URIs, and storing the
   * contents in a provided concrete JSON object.
   */
  template <Adapter A, typename... Args>
  Schema(A const & json, schema::Version version, URIResolver<A> resolve, Args &&... args)
      : Schema(json, version, DocumentCache<A>(resolve), std::forward<Args>(args)...) {}

  /**
   * @param json An adapter to a json schema
   *
   * @param version The json-schema draft version that all schemas will prefer
   */
  template <Adapter A, Not<DocumentCache<A>>... Args>
  Schema(A const & json, schema::Version version, Args &&... args)
      : Schema(json, version, DocumentCache<A>(), std::forward<Args>(args)...) {}

  /**
   * @param json Any non-adapter (JSON) object. Will be immedately converted
   * into an Adapter object to allow us to walk through it w/o specialization.
   */
  template <typename JSON, typename... Args>
    requires(not Adapter<JSON>)
  explicit Schema(JSON const & json, Args &&... args)
      : Schema(adapter::AdapterFor<JSON const>(json), std::forward<Args>(args)...) {}

private:
  /**
   * @brief Cache an alias to a given schema, without ownership. alias_cache_ is
   * a many-to-one association.
   * Syntactic sugar for "add pointer to map and return".
   *
   * @param where The key aliasing the schema, which may also be the original
   * lexical key.
   *
   * @param schema The pointer to a schema being stored
   */
  schema::Node const * alias(detail::Reference const & where, schema::Node const * schema) {
    alias_cache_.emplace(where, schema);
    return schema;
  }

  /**
   * @brief Syntactic sugar for finding a map value as an optional instead of an
   * iterator that may be "end".
   *
   * @param ref The key being looked up
   */
  std::optional<schema::Node const *> from_cache(detail::Reference const & ref) {
    if (auto it = alias_cache_.find(ref); it != alias_cache_.end()) {
      return it->second;
    }

    return std::nullopt;
  }

  /**
   * @brief Resolve a $ref/$dynamicRef tag and construct or load from cache the
   * schema that is being pointed to.
   *
   * @param context All of the context information about the schema, importantly
   * the location information, {@see jvalidate::detail::ReferenceManager}, and
   * {@see jvalidate::detail::Vocabulary}.
   *
   * @param dynamic_reference Is this request coming from a "$dynamicRef"/
   * "$recursiveRef" tag, or a regular "$ref" tag.
   *
   * @returns A schema node, that will also be stored in a local cache.
   *
   * @throws std::runtime_error if the reference is to an unloaded URI, and we
   * fail to load it. If the preprocessor definition
   * JVALIDATE_LOAD_FAILURE_AS_FALSE_SCHEMA is set, then we instead return an
   * always-false schema with a custom error message. This is primarily for use
   * in writing tests for JSON-Schema's selfvalidation test cases.
   */
  template <Adapter A>
  schema::Node const * resolve(detail::Reference const & ref,
                               detail::ParserContext<A> const & context, bool dynamic_reference) {
    detail::Reference lexical = context.ref.canonicalize(ref, context.where, dynamic_reference);
    detail::Reference dynamic = dynamic_reference ? lexical : context.dynamic_where / "$ref";

    if (std::optional cached = from_cache(dynamic)) {
      return *cached;
    }

    if (std::optional root = context.ref.load(lexical, context.vocab)) {
      return fetch_schema(context.rebind(*root, lexical, dynamic));
    }

    std::string error = "URIResolver could not resolve " + std::string(lexical.uri());
#ifdef JVALIDATE_LOAD_FAILURE_AS_FALSE_SCHEMA
    return alias(dynamic, &cache_.try_emplace(dynamic, error).first->second);
#else
    JVALIDATE_THROW(std::runtime_error, error);
#endif
  }

  /**
   * @brief Fetch from cache or create a new schema node from the given context,
   * which may be the result of resolving a reference {@see Schema::resolve}, or
   * simply loading a child schema via {@see ParserContext::node}.
   *
   * @param context The current operating context of the schema
   */
  template <Adapter A> schema::Node const * fetch_schema(detail::ParserContext<A> const & context) {
    // TODO(samjaffe): No longer promises uniqueness - instead track unique URI's
    if (std::optional cached = from_cache(context.dynamic_where)) {
      return *cached;
    }

    adapter::Type const type = context.schema.type();
    // Boolean schemas were made universally permitted in Draft06. Before then,
    // you could only use them for specific keywords, like additionalProperties.
    if (type == adapter::Type::Boolean && context.vocab->version() >= schema::Version::Draft06) {
      return alias(context.dynamic_where, context.schema.as_boolean() ? &accept_ : &reject_);
    }

    // If the schema is not universal accept/reject, then it MUST be an object
    EXPECT_M(type == adapter::Type::Object, "invalid schema at " << context.dynamic_where);
    // The empty object is equivalent to true, but is permitted in prior drafts
    if (context.schema.object_size() == 0) {
      return alias(context.dynamic_where, &accept_);
    }

    // Because of the below alias() expression, and the above from_cache
    // expression, it shouldn't be possible for try_emplace to not create a new
    // schema node. We keep the check in anyway just in case somehow things have
    // gotten into a malformed state.
    auto [it, created] = cache_.try_emplace(context.dynamic_where);
    EXPECT_M(created, "creating duplicate schema at... " << context.dynamic_where);

    // Do this here first in order to protect from infinite loops
    alias(context.dynamic_where, &it->second);
    it->second.construct(context);
    return &it->second;
  }
};
}

namespace jvalidate::detail {
template <Adapter A> schema::Node const * ParserContext<A>::node() const {
  return root.fetch_schema(*this);
}

template <Adapter A> schema::Node const * ParserContext<A>::always() const {
  return fixed_schema(schema.as_boolean());
}

template <Adapter A> schema::Node const * ParserContext<A>::fixed_schema(bool accept) const {
  return accept ? &root.accept_ : &root.reject_;
}
}

namespace jvalidate::schema {
template <Adapter A>
detail::OnBlockExit Node::resolve_anchor(detail::ParserContext<A> const & context) {
  auto const schema = context.schema.as_object();

  if (context.vocab->version() < schema::Version::Draft2019_09 || not schema.contains("$id")) {
    return nullptr;
  }

  return context.ref.dynamic_scope(context.where);
}

template <Adapter A> bool Node::resolve_reference(detail::ParserContext<A> const & context) {
  auto const schema = context.schema.as_object();

  if (schema.contains("$ref")) {
    detail::Reference ref(schema["$ref"].as_string());

    reference_ = context.root.resolve(ref, context, false);
    return true;
  }

  // Prior to Draft2019-09, "$ref" was the only way to reference another
  // schema (ignoring Draft03's extends keyword, which was more like allOf)
  if (context.vocab->version() < Version::Draft2019_09) {
    return false;
  }

  std::string const dyn_ref =
      context.vocab->version() > schema::Version::Draft2019_09 ? "$dynamicRef" : "$recursiveRef";
  if (schema.contains(dyn_ref)) {
    detail::Reference ref(schema[dyn_ref].as_string());

    reference_ = context.root.resolve(ref, context, true);
    return true;
  }

  return false;
}

template <Adapter A> void Node::construct(detail::ParserContext<A> context) {
  EXPECT(context.schema.type() == adapter::Type::Object);

  auto const schema = context.schema.as_object();

  if (schema.contains("$schema")) {
    // At any point in the schema, we're allowed to change versions
    // This means that we're not version-locked to the latest grammar
    // (which is especially important for some breaking changes)
    context.vocab = &context.ref.vocab(URI(schema["$schema"].as_string()));
  }

  auto _ = resolve_anchor(context);
  bool const has_reference = resolve_reference(context);

  if (schema.contains("default")) {
    default_ = schema["default"].freeze();
  }

  if (schema.contains("description")) {
    description_ = schema["description"].as_string();
  }

  // Prior to Draft 2019-09, reference keywords take precedence over everything
  // else (instead of allowing direct extensions).
  if (has_reference && context.vocab->version() < Version::Draft2019_09) {
    return;
  }

  for (auto const & [key, subschema] : schema) {
    // Using a constraint store allows overriding certain rules, or the creation
    // of user-defined extention vocabularies.
    if (not context.vocab->is_constraint(key)) {
      continue;
    }
    // A constraint may return null if it is not applicable - but otherwise
    // well-formed. For example, before Draft-06 "exclusiveMaximum" was a
    // modifier property for "maximum", and not a unique constaint on its own.
    // Therefore, we parse it alongside parsing "maximum", and could return
    // nullptr when requesting a constraint pointer for "exclusiveMaximum".
    auto [constraint, post] = context.vocab->constraint(key, context.child(subschema, key));
    if (not constraint) {
      continue;
    }
    if (post) {
      post_constraints_.emplace(key, std::move(constraint));
    } else {
      constraints_.emplace(key, std::move(constraint));
    }
  }
}
}