#pragma once

#include <functional>
#include <map>
#include <set>
#include <unordered_map>

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

namespace jvalidate {
template <Adapter A> class ReferenceManager {
public:
  using Keywords = std::unordered_map<std::string_view, std::set<schema::Wraps>>;

private:
  DocumentCache<A> & external_;
  detail::ReferenceCache references_;

  std::map<detail::RootReference, A> roots_;

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

public:
  ReferenceManager(DocumentCache<A> & external, A const & root, schema::Version version,
                   Keywords const & keywords)
      : external_(external), roots_{{{}, root}} {
    prime(root, {}, version, keywords);
  }

  detail::OnBlockExit suppress(detail::Anchor const & anchor) {
    if (auto it = active_dynamic_anchors_.find(anchor); it != active_dynamic_anchors_.end()) {
      detail::Reference where = it->second;
      active_dynamic_anchors_.erase(it);
      return [this, anchor, where]() { active_dynamic_anchors_[anchor] = where; };
    }
    return nullptr;
  }

  auto scoped_activate(detail::Reference ref) {
    URI const uri = references_.relative_to_nearest_anchor(ref).uri();
    activate_dynamic(uri);
    return [this, uri]() { deactivate_dynamic(uri); };
  }

  auto load(detail::Reference const & ref, detail::ParserContext<A> const & context)
      -> std::pair<std::optional<A>, detail::OnBlockExit> {
    auto scoped_schema = [this, &ref](A const & schema) {
      activate_dynamic(ref.uri());
      return std::make_pair(ref.pointer().walk(schema),
                            [this, uri = ref.uri()] { deactivate_dynamic(uri); });
    };

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

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

    // TODO(samjaffe): Change Versions if needed...
    references_.emplace(ref.uri());
    prime(*external, ref, context.version, context.factory.keywords(context.version));

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

    // Will get called if the external schema does not declare a root document id?
    return scoped_schema(*external);
  }

  detail::Reference canonicalize(detail::Reference const & ref, detail::Reference const & parent,
                                 detail::inout<bool> dynamic_reference) const {
    if (auto it = active_dynamic_anchors_.find(ref.anchor()); it != active_dynamic_anchors_.end()) {
      if (dynamic_reference) {
        return it->second;
      }
      dynamic_reference = ref.uri().empty() && ref.pointer().empty();
    }

    if (active_dynamic_anchors_.empty()) {
      dynamic_reference = true;
    }

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

    URI uri = [this, &ref, &parent]() {
      if (ref.uri().empty() && parent.uri().empty()) {
        return references_.actual_parent_uri(parent);
      }

      URI uri = ref.uri().empty() ? parent.uri() : ref.uri();
      if (not uri.is_rootless()) {
        return uri;
      }

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

    detail::Reference rval(uri, ref.anchor(), ref.pointer());

    // Will now need to go make an external fetch...
    // TODO(samjaffe): Make that process internal?
    return rval;
  }

private:
  void activate_dynamic(URI const & uri) {
    active_dynamic_anchors_.insert(dynamic_anchors_[uri].begin(), dynamic_anchors_[uri].end());
  }

  void deactivate_dynamic(URI const & uri) {
    std::map<detail::Anchor, detail::Reference> const & candidates = dynamic_anchors_[uri];
    for (auto it = active_dynamic_anchors_.begin(); it != active_dynamic_anchors_.end();) {
      if (auto cit = candidates.find(it->first); it->second == cit->second) {
        it = active_dynamic_anchors_.erase(it);
      } else {
        ++it;
      }
    }
  }

  void prime(Adapter auto const & json, detail::Reference where, schema::Version version,
             Keywords const & keywords) {
    if (json.type() != adapter::Type::Object) {
      return;
    }

    canonicalize(where, version, json);

    for (auto const & [key, value] : json.as_object()) {
      auto vit = keywords.find(key);
      if (vit == keywords.end()) {
        continue;
      }

      if (vit->second.contains(schema::Wraps::Array) && value.type() == adapter::Type::Array) {
        size_t index = 0;
        for (auto const & elem : value.as_array()) {
          prime(elem, where / key / index, version, keywords);
          ++index;
        }
      } else if (vit->second.contains(schema::Wraps::Object) &&
                 value.type() == adapter::Type::Object) {
        for (auto const & [prop, elem] : value.as_object()) {
          prime(elem, where / key / prop, version, keywords);
        }
      } else if (vit->second.contains(schema::Wraps::Schema)) {
        prime(value, where / key, version, keywords);
      }
    }
  }

  void canonicalize(detail::Reference & where, schema::Version version, A const & json) {
    std::string const id = version <= schema::Version::Draft04 ? "id" : "$id";
    auto const schema = json.as_object();

    detail::RootReference root = where.root();
    if (schema.contains(id)) {
      root = detail::RootReference(schema[id].as_string());
      if (root.uri().empty()) {
        root = detail::RootReference(where.uri(), root.anchor());
      } else if (not root.uri().is_rootless() || where.uri().empty()) {
        // By definition - rootless URIs cannot be relative
      } else if (root.uri().is_relative()) {
        root = detail::RootReference(where.uri().parent() / root.uri(), root.anchor());
      } else {
        root = detail::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 = detail::RootReference(root.uri(), detail::Anchor(schema["$anchor"].as_string()));
      roots_.emplace(root, json);
      where = references_.emplace(where, root);
    }

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

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

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

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

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

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