#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/document_cache.h>
#include <jvalidate/enum.h>
#include <jvalidate/forward.h>

namespace jvalidate::schema {
class Node {
private:
  std::string description_;
  std::unique_ptr<adapter::Const const> default_{nullptr};

  detail::Reference uri_;
  std::optional<std::string> rejects_all_;
  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_{};

protected:
  static Version schema_version(std::string_view url);
  static Version schema_version(Adapter auto const & json);
  static Version schema_version(Adapter auto const & json, Version default_version);

public:
  Node() = default;
  Node(std::string const & rejection_reason) : rejects_all_(rejection_reason) {}
  template <Adapter A> void construct(detail::ParserContext<A> context);

  bool is_pure_reference() const {
    return reference_ && constraints_.empty() && post_constraints_.empty() && not default_;
  }
  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 constraints_; }

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

private:
  template <Adapter A> detail::OnBlockExit resolve_anchor(detail::ParserContext<A> & context);
  template <Adapter A> bool resolve_reference(detail::ParserContext<A> const & context);
};

inline Version Node::schema_version(std::string_view url) {
  static std::map<std::string_view, Version> const g_schema_ids{
      {"json-schema.org/draft-04/schema", Version::Draft04},
      {"json-schema.org/draft-06/schema", Version::Draft06},
      {"json-schema.org/draft-07/schema", Version::Draft07},
      {"json-schema.org/draft/2019-09/schema", Version::Draft2019_09},
      {"json-schema.org/draft/2020-12/schema", Version::Draft2020_12},
  };

  if (url.ends_with('#')) {
    url.remove_suffix(1);
  }
  if (url.starts_with("http://") || url.starts_with("https://")) {
    url.remove_prefix(url.find(':') + 3);
  }

  auto it = g_schema_ids.find(url);
  EXPECT_T(it != g_schema_ids.end(), std::invalid_argument, url);
  return it->second;
}

Version Node::schema_version(Adapter auto const & json) {
  EXPECT(json.type() == adapter::Type::Object);
  EXPECT(json.as_object().contains("$schema"));

  auto const & schema = json.as_object()["$schema"];
  EXPECT(schema.type() == adapter::Type::String);

  return schema_version(schema.as_string());
}

Version Node::schema_version(Adapter auto const & json, Version default_version) {
  RETURN_UNLESS(json.type() == adapter::Type::Object, default_version);
  RETURN_UNLESS(json.as_object().contains("$schema"), default_version);

  auto const & schema = json.as_object()["$schema"];
  RETURN_UNLESS(schema.type() == adapter::Type::String, default_version);

  return schema_version(schema.as_string());
}
}

namespace jvalidate {

class Schema : public schema::Node {
private:
  friend class schema::Node;
  template <Adapter A> friend class detail::ParserContext;
  struct DynamicRef {
    template <typename F>
    DynamicRef(detail::Reference const & where, F const & reconstruct)
        : where(where), reconstruct(reconstruct) {}

    detail::Reference where;
    std::function<schema::Node const *()> reconstruct;
  };

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

  // A map of (URI, Anchor) => (URI, Pointer), binding an anchor reference
  // to it's fully resolved path.
  std::map<detail::Reference, detail::Reference> anchors_;

  // A map of anchors to DynamicRef info - note that DynamicRef.reconstruct is
  // an unsafe object, because it holds an object which may hold references to
  // temporary objects.
  // Nothing should be added to this object except through calling
  // {@see Node::resolve_anchor}, which returns a scope(exit) construct that
  // cleans up the element.
  std::map<detail::Anchor, DynamicRef> dynamic_anchors_;

  // 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;
    }

    external.cache_reference(URI(), json);
    detail::ParserContext<A> root{*this, json, version, factory, external};
    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 An adapter to a json schema
   */
  template <Adapter A, Not<schema::Version>... Args>
  explicit Schema(A const & json, Args &&... args)
      : Schema(json, schema_version(json), 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>
  explicit Schema(JSON const & json, Args &&... args)
      : Schema(adapter::AdapterFor<JSON const>(json), std::forward<Args>(args)...) {}

private:
  /**
   * @brief Associate an anchor with its absolute path
   * @pre We should not already have an anchor associated with this anchor
   *
   * @param anchor A URI-Reference containing only a URI and Anchor
   * @param from A URI-Reference representing the absolute path to this Anchor
   */
  void anchor(detail::Reference const & anchor, detail::Reference const & from) {
    EXPECT_M(anchors_.try_emplace(anchor.root(), from).second,
             "more than one anchor found for uri " << anchor);
  }

  template <Adapter A>
  void dynamic_anchor(detail::Anchor const & anchor, detail::ParserContext<A> const & context) {
    dynamic_anchors_.try_emplace(anchor, context.where,
                                 [this, context]() { return fetch_schema(context); });
  }

  void remove_dynamic_anchor(detail::Anchor const & anchor, detail::Reference const & where) {
    if (auto it = dynamic_anchors_.find(anchor);
        it != dynamic_anchors_.end() && it->second.where == where) {
      dynamic_anchors_.erase(it);
    }
  }

  schema::Node const * alias(detail::Reference const & where, schema::Node const * schema) {
    EXPECT_M(alias_cache_.try_emplace(where, schema).second,
             "more than one schema found with uri " << where);
    return schema;
  }

  std::optional<schema::Node const *> from_cache(detail::Reference ref) {
    if (auto it = anchors_.find(ref.root()); it != anchors_.end()) {
      ref = it->second / ref.pointer();
    }

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

    return std::nullopt;
  }

  template <Adapter A>
  schema::Node const * resolve(detail::Reference ref, detail::ParserContext<A> const & context) {
    // Special case if the root-level document does not have an $id property
    if (ref == detail::Reference() && context.where.uri().empty()) {
      return this;
    }
    if (ref.uri().empty()) {
      ref = detail::Reference(context.where.uri(), ref.anchor(), ref.pointer());
    }

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

    // SPECIAL RULE: Resolve this URI into the context of the calling URI
    if (not ref.uri().empty() && ref.uri().scheme().empty()) {
      URI const & relative_to = context.where.uri();
      EXPECT_M(relative_to.resource().rfind('/') != std::string::npos,
               "Unable to deduce root for relative uri " << ref.uri() << " (" << relative_to
                                                         << ")");
      ref = detail::Reference(relative_to.parent() / ref.uri(), ref.anchor(), ref.pointer());
    }

    std::optional schema = context.external.try_load(ref);
    if (not schema.has_value()) {
      std::string error = "URIResolver could not resolve " + std::string(ref.uri());
      return alias(ref, &cache_.try_emplace(ref, error).first->second);
    }

    return fetch_schema(context.rebind(*schema, ref));
  }

  schema::Node const * resolve_dynamic(detail::Anchor const & ref) {
    auto it = dynamic_anchors_.find(ref);
    EXPECT_M(it != dynamic_anchors_.end(), "Unmatched $dynamicRef '" << ref << "'");
    return it->second.reconstruct();
  }

  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.where)) {
      return *cached;
    }

    adapter::Type const type = context.schema.type();
    if (type == adapter::Type::Boolean && context.version >= schema::Version::Draft06) {
      return alias(context.where, context.schema.as_boolean() ? &accept_ : &reject_);
    }

    EXPECT_M(type == adapter::Type::Object, "invalid schema at " << context.where);
    if (context.schema.object_size() == 0) {
      return alias(context.where, &accept_);
    }

    auto [it, created] = cache_.try_emplace(context.where);
    EXPECT_M(created, "creating duplicate schema at... " << context.where);

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

    // Special Case - if the only is the reference constraint, then we don't need
    // to store it uniquely. Draft2019_09 supports directly extending a $ref schema
    // in the same schema, instead of requiring an allOf clause.
    schema::Node const * node = *it->second.reference_schema();
    cache_.erase(it);
    return alias_cache_[context.where] = node;
  }
};
}

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> & context) {
  auto const schema = context.schema.as_object();

  if (schema.contains("$anchor")) {
    // Create an anchor mapping using the current document and the anchor
    // string. There's no need for special validation/chaining here, because
    // {@see Schema::resolve} will turn all $ref/$dynamicRef anchors into
    // their fully-qualified path.
    detail::Anchor anchor(schema["$anchor"].as_string());
    context.root.anchor(detail::Reference(context.where.uri(), anchor), context.where);
    return nullptr;
  }

  if (context.version == Version::Draft2019_09 && schema.contains("$recursiveAnchor")) {
    EXPECT_M(schema["$recursiveAnchor"].as_boolean(), "$recursiveAnchor MUST be 'true'");

    context.root.dynamic_anchor(detail::Anchor(), context);
    return [&context]() { context.root.remove_dynamic_anchor(detail::Anchor(), context.where); };
  }

  if (context.version > Version::Draft2019_09 && schema.contains("$dynamicAnchor")) {
    detail::Anchor anchor(schema["$dynamicAnchor"].as_string());

    context.root.dynamic_anchor(anchor, context);
    return [&context, anchor]() { context.root.remove_dynamic_anchor(anchor, context.where); };
  }

  return nullptr;
}

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);
    return true;
  }

  if (context.version < Version::Draft2019_09) {
    return false;
  }

  if (context.version == Version::Draft2019_09 && schema.contains("$recursiveRef")) {
    detail::Reference ref(schema["$recursiveRef"].as_string());
    EXPECT_M(ref == detail::Reference(), "Only the root schema is permitted as a $recursiveRef");

    reference_ = context.root.resolve_dynamic(detail::Anchor());
    return true;
  }

  if (context.version > Version::Draft2019_09 && schema.contains("$dynamicRef")) {
    detail::Reference ref(schema["$dynamicRef"].as_string());

    reference_ = context.root.resolve_dynamic(ref.anchor());
    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.version = schema_version(context.schema);
  }

  if (schema.contains("$id")) {
    detail::Reference id(schema["$id"].as_string(), false);
    if (id.uri().scheme().empty() and not context.where.uri().empty()) {
      id = detail::Reference(context.where.uri().parent() / id.uri(), {}, id.pointer());
    }

    if (id != context.where) {
      context.external.cache_reference(id.uri(), context.schema);
      context.root.alias(context.where = id, this);
    }
  }

  [[maybe_unused]] 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.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.
    auto make_constraint = context.factory(key, context.version);
    if (not make_constraint) {
      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 = make_constraint(context.child(subschema, key));
    if (not constraint) {
      continue;
    }
    if (context.factory.is_post_constraint(key)) {
      post_constraints_.emplace(key, std::move(constraint));
    } else {
      constraints_.emplace(key, std::move(constraint));
    }
  }
}
}