#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>
#include <jvalidate/reference_handler.h>

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

  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_;
  }
  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:
  template <Adapter A> detail::OnBlockExit resolve_anchor(detail::ParserContext<A> const & 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;

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

    ReferenceManager<A> ref(external, json, version, factory.keywords(version));
    detail::ParserContext<A> root{*this, json, version, factory, 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 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:
  schema::Node const * alias(detail::Reference const & where, schema::Node const * schema) {
    alias_cache_.emplace(where, schema);
    return schema;
  }

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

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

    // Special case if the root-level document does not have an $id property
    if (auto [root, scope] = context.ref.load(lexical, context); root.has_value()) {
      return fetch_schema(context.rebind(*root, lexical, dynamic));
    }

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

  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();
    if (type == adapter::Type::Boolean && context.version >= schema::Version::Draft06) {
      return alias(context.dynamic_where, context.schema.as_boolean() ? &accept_ : &reject_);
    }

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

    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 (schema.contains("$anchor") || context.version < schema::Version::Draft2019_09) {
    return nullptr;
  }

  if (context.version > schema::Version::Draft2019_09) {
    return context.ref.scoped_activate(context.where);
  }

  if (context.version != schema::Version::Draft2019_09) {
    return nullptr;
  }
  if (schema.contains("$recursiveAnchor") && schema["$recursiveAnchor"].as_boolean()) {
    return context.ref.scoped_activate(context.where);
  }
  if (not schema.contains("$id") && not schema.contains("$recursiveAnchor")) {
    return nullptr;
  }

  return context.ref.suppress(detail::Anchor());
}

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

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

  std::string const dyn_ref =
      context.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.version = schema_version(context.schema);
  }

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