#pragma once

#include <memory>
#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/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_;
  bool rejects_all_{false};
  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(bool rejects_all = false) : rejects_all_(rejects_all) {}
  template <Adapter A> Node(detail::ParserContext<A> context);

  bool is_pure_reference() const {
    return reference_ && constraints_.empty() && post_constraints_.empty() && not default_;
  }
  bool 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{
      {"http://json-schema.org/draft-04/schema", Version::Draft04},
      {"http://json-schema.org/draft-06/schema", Version::Draft06},
      {"http://json-schema.org/draft-07/schema", Version::Draft07},
      {"http://json-schema.org/draft/2019-09/schema", Version::Draft2019_09},
      {"http://json-schema.org/draft/2020-12/schema", Version::Draft2020_12},
  };

  if (url.ends_with('#')) {
    url.remove_suffix(1);
  }

  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_{true};
  schema::Node reject_{false};

  std::map<detail::Reference, detail::Reference> anchors_;
  std::map<detail::Anchor, DynamicRef> dynamic_anchors_;
  std::map<detail::Reference, schema::Node> cache_;

  std::map<detail::Reference, schema::Node const *> alias_cache_;

public:
  template <Adapter A>
  explicit Schema(A const & json, ConstraintFactory<A> const & factory = {})
      : Schema(json, schema_version(json), factory) {}

  template <Adapter A>
  Schema(A const & json, schema::Version version, ConstraintFactory<A> const & factory = {})
      : schema::Node(detail::ParserContext<A>{
            .root = *this, .schema = json, .version = version, .factory = factory}) {}

  template <typename JSON, typename... Args>
  explicit Schema(JSON const & json, Args &&... args)
      : Schema(adapter::AdapterFor<JSON const>(json), std::forward<Args>(args)...) {}

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

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

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

  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) {
    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(type == adapter::Type::Object);
    if (context.schema.object_size() == 0) {
      return alias(context.where, &accept_);
    }

    auto [it, created] = cache_.try_emplace(context.where, context);
    EXPECT_M(created, "more than one schema found with uri " << context.where);

    schema::Node const * node = &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.
    if (node->is_pure_reference()) {
      node = *node->reference_schema();
      cache_.erase(it);
      return alias(context.where, node);
    }

    return alias(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 schema.as_boolean() ? &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); };
  }
}

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.template resolve<A>(ref, context.where, context.version);
    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> Node::Node(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")) {
    context.root.alias(detail::Reference(schema["$id"].as_string(), false), 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();
  }

  for (auto [key, subschema] : schema) {
    // Using a constraint store allows overriding certain rules, or the creation
    // of user-defined extention vocabularies.
    if (auto make_constraint = context.factory(key, context.version)) {
      EXPECT_M(not has_reference || context.version >= Version::Draft2019_09,
               "Cannot directly extend $ref schemas before Draft2019-09");
      // 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".
      if (auto constraint = make_constraint(context.child(subschema, key))) {
        auto & into = context.factory.is_post_constraint(key) ? post_constraints_ : constraints_;
        into.emplace(key, std::move(constraint));
      }
    }
  }
}
}