#pragma once

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

#include <jvalidate/detail/anchor.h>
#include <jvalidate/detail/pointer.h>
#include <jvalidate/detail/reference.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:
  std::map<detail::RootReference, A> roots_;

  std::map<detail::Reference, detail::RootReference, std::greater<>> absolute_to_canonical_;
  std::map<detail::RootReference, detail::Reference, std::greater<>> canonical_to_absolute_;

public:
  ReferenceManager(A const & root, schema::Version version, Keywords const & keywords)
      : roots_{{{}, root}} {
    prime_impl(root, {}, version, keywords);
  }

  std::optional<A> root(detail::RootReference const & root) const {
    if (auto it = roots_.find(root); it != roots_.end()) {
      return it->second;
    }
    return std::nullopt;
  }

  detail::Reference canonicalize(detail::Reference const & ref) const {
    if (canonical_to_absolute_.contains(ref.root())) {
      return ref;
    }

    auto it = absolute_to_canonical_.upper_bound(ref);
    if (it == absolute_to_canonical_.end()) {
      return ref;
    }

    auto const & [absolute, anchor] = *it;
    if (not ref.pointer().starts_with(absolute.pointer())) {
      return ref;
    }

    return detail::Reference(anchor, ref.pointer().relative_to(absolute.pointer()));
  }

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

    // TODO(samjaffe): Clean this clause up
    URI uri = ref.uri().empty() ? parent.root().uri() : ref.uri();
    if (uri.empty()) {
      auto parent_uri = canonicalize(parent).uri();
      if (parent_uri == parent.uri() && parent.uri().empty()) {
        auto it = absolute_to_canonical_.find(detail::Reference());
        if (it != absolute_to_canonical_.end()) {
          uri = it->second.uri();
        }
      }
    } else if (uri.is_rootless()) {
      auto parent_uri = canonicalize(parent).uri();
      if (parent_uri == parent.uri() && parent.uri().empty()) {
        auto it = absolute_to_canonical_.find(detail::Reference());
        if (it != absolute_to_canonical_.end()) {
          parent_uri = it->second.uri();
        }
      }
      EXPECT_M(parent_uri.resource().rfind('/') != std::string::npos,
               "Unable to deduce root for relative uri " << uri << " (" << parent_uri << ")");
      uri = (uri.is_relative() ? parent_uri.parent() : parent_uri.root()) / uri;
    }

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

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

  void prime(Adapter auto const & json, detail::RootReference const & where,
             schema::Version version, Keywords const & keywords) {
    canonical_to_absolute_.emplace(where, detail::Reference(where));
    prime_impl(json, detail::Reference(where), version, keywords);
  }

private:
  void prime_impl(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_impl(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_impl(elem, where / key / prop, version, keywords);
        }
      } else if (vit->second.contains(schema::Wraps::Schema)) {
        prime_impl(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 (root.uri().is_relative() && not where.uri().empty()) {
        root = detail::RootReference(where.uri().parent() / root.uri(), root.anchor());
      } else if (root.uri().is_rootless() && not where.uri().empty()) {
        root = detail::RootReference(where.uri().root() / root.uri(), root.anchor());
      }

      cache(root, where, json);
    }

    if (not schema.contains("$anchor") || version < schema::Version::Draft2019_09) {
      return;
    }

    EXPECT_M(root.anchor().empty(), "Cannot have $id with anchor and $anchor tags at same time");
    root = detail::RootReference(root.uri(), detail::Anchor(schema["$anchor"].as_string()));

    cache(root, where, json);
  }

  template <size_t I>
  bool extract_relative(detail::Reference const & where, auto const & tup,
                        detail::Pointer & out) const {
    auto const & ptr = std::get<I>(tup).pointer();
    if (where.pointer().starts_with(ptr)) {
      out = where.pointer().relative_to(ptr);
      return true;
    }
    return false;
  }

  void recursive_cache(detail::RootReference const & root, detail::Reference const & where) {
    if (where.pointer().empty()) {
      return;
    }
    detail::Pointer relative;
    if (auto it = canonical_to_absolute_.lower_bound(where.root());
        it != canonical_to_absolute_.end() && where.root() != it->second.root() &&
        extract_relative<1>(where, *it, relative)) {
      absolute_to_canonical_.emplace(it->second / relative, root);
      recursive_cache(root, it->second / relative);
    }
  }

  void cache(detail::RootReference const & root, detail::Reference & where,
             Adapter auto const & json) {
    recursive_cache(root, where);
    absolute_to_canonical_.emplace(where, root);
    canonical_to_absolute_.emplace(root, where);

    roots_.emplace(root, json);
    where = detail::Reference(root);
  }
};
}