// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#ifndef KOKKOS_HOST_GRAPH_IMPL_HPP
#define KOKKOS_HOST_GRAPH_IMPL_HPP

#include <Kokkos_ExecPolicy.hpp>
#include <Kokkos_Graph.hpp>

#include <impl/Kokkos_GraphImpl_fwd.hpp>
#include <impl/Kokkos_Default_Graph_fwd.hpp>

#include <Serial/Kokkos_Serial.hpp>
#include <OpenMP/Kokkos_OpenMP.hpp>
// FIXME @graph other backends?

#include <impl/Kokkos_EBO.hpp>

#include <set>

namespace Kokkos {
namespace Impl {

//==============================================================================
// <editor-fold desc="GraphImpl default implementation"> {{{1

template <class ExecutionSpace>
struct GraphImpl : private InstanceStorage<ExecutionSpace> {
 public:
  using root_node_impl_t =
      GraphNodeImpl<ExecutionSpace, Kokkos::Experimental::TypeErasedTag,
                    Kokkos::Experimental::TypeErasedTag>;

  using aggregate_impl_t = GraphNodeAggregateDefaultImpl<ExecutionSpace>;

 private:
  using execution_space_instance_storage_base_t =
      InstanceStorage<ExecutionSpace>;

  using node_details_t = GraphNodeBackendSpecificDetails<ExecutionSpace>;
  std::set<std::shared_ptr<node_details_t>> m_sinks;

 public:
  //----------------------------------------------------------------------------
  // <editor-fold desc="Constructors, destructor, and assignment"> {{{2

  // Not movable or copyable; it spends its whole live as a shared_ptr in the
  // Graph object
  GraphImpl()                            = default;
  GraphImpl(GraphImpl const&)            = delete;
  GraphImpl(GraphImpl&&)                 = delete;
  GraphImpl& operator=(GraphImpl const&) = delete;
  GraphImpl& operator=(GraphImpl&&)      = delete;
  ~GraphImpl()                           = default;

  explicit GraphImpl(ExecutionSpace arg_space)
      : execution_space_instance_storage_base_t(std::move(arg_space)) {}

  // </editor-fold> end Constructors, destructor, and assignment }}}2
  //----------------------------------------------------------------------------

  ExecutionSpace const& get_execution_space() const {
    return this->execution_space_instance_storage_base_t::instance();
  }

  //----------------------------------------------------------------------------
  // <editor-fold desc="required customizations"> {{{2

  template <class NodeImpl>
  void add_node(std::shared_ptr<NodeImpl> arg_node_ptr) {
    static_assert(
        Kokkos::Impl::is_specialization_of_v<NodeImpl, GraphNodeImpl>);
    // Since this is always called before any calls to add_predecessor involving
    // it, we can treat this node as a sink until we discover otherwise.
    arg_node_ptr->node_details_t::set_kernel(arg_node_ptr->get_kernel());
    auto spot = m_sinks.find(arg_node_ptr);
    KOKKOS_ASSERT(spot == m_sinks.end())
    m_sinks.insert(std::move(spot), std::move(arg_node_ptr));
  }

  template <class NodeImplPtr, class PredecessorRef>
  // requires PredecessorRef is a specialization of GraphNodeRef that has
  // already been added to this graph and NodeImpl is a specialization of
  // GraphNodeImpl that has already been added to this graph.
  void add_predecessor(NodeImplPtr arg_node_ptr, PredecessorRef arg_pred_ref) {
    auto node_ptr_spot = m_sinks.find(arg_node_ptr);
    auto pred_ptr      = GraphAccess::get_node_ptr(arg_pred_ref);
    auto pred_ref_spot = m_sinks.find(pred_ptr);
    KOKKOS_ASSERT(node_ptr_spot != m_sinks.end())
    if (pred_ref_spot != m_sinks.end()) {
      // delegate responsibility for executing the predecessor to arg_node
      // and then remove the predecessor from the set of sinks
      (*node_ptr_spot)->set_predecessor(std::move(*pred_ref_spot));
      m_sinks.erase(pred_ref_spot);
    } else {
      // We still want to check that it's executed, even though someone else
      // should have executed it before us
      (*node_ptr_spot)->set_predecessor(std::move(pred_ptr));
    }
  }

  template <class... PredecessorRefs>
  // See requirements/expectations in GraphBuilder
  auto create_aggregate_ptr(PredecessorRefs&&...) {
    // The attachment to predecessors, which is all we really need, happens
    // in the generic layer, which calls through to add_predecessor for
    // each predecessor ref, so all we need to do here is create the (trivial)
    // aggregate node.
    using aggregate_node_impl_t =
        GraphNodeImpl<ExecutionSpace, aggregate_impl_t,
                      Kokkos::Experimental::TypeErasedTag>;
    return GraphAccess::make_node_shared_ptr<aggregate_node_impl_t>(
        this->get_execution_space(), _graph_node_kernel_ctor_tag{},
        aggregate_impl_t{});
  }

  auto create_root_node_ptr() {
    auto rv = Kokkos::Impl::GraphAccess::make_node_shared_ptr<root_node_impl_t>(
        get_execution_space(), _graph_node_is_root_ctor_tag{});
    m_sinks.insert(rv);
    return rv;
  }

  void instantiate() {
    KOKKOS_EXPECTS(!m_has_been_instantiated);
    m_has_been_instantiated = true;
  }

  void submit(const ExecutionSpace& exec) {
    if (!m_has_been_instantiated) instantiate();
    // This reset is gross, but for the purposes of our simple host
    // implementation...
    for (auto& sink : m_sinks) {
      sink->reset_has_executed();
    }

    // We don't know where the nodes will execute, so we need to fence the given
    // execution space instance before proceeding. This is the simplest way
    // of guaranteeing that the kernels in the graph are correctly "enqueued".
    exec.fence(
        "Kokkos::DefaultGraph::submit: fencing before launching graph nodes");

    for (auto& sink : m_sinks) {
      sink->execute_node(exec);
    }

    // Once all sinks have been executed, we need to fence them.
    for (const auto& sink : m_sinks) {
      if (sink->awaitable() && sink->get_execution_space() != exec)
        sink->get_execution_space().fence(
            "Kokkos::DefaultGraph::submit: fencing before ending graph submit");
    }
  }

 private:
  bool m_has_been_instantiated = false;

  // </editor-fold> end required customizations }}}2
  //----------------------------------------------------------------------------
};

// </editor-fold> end GraphImpl default implementation }}}1
//==============================================================================

}  // end namespace Impl

}  // end namespace Kokkos

#include <impl/Kokkos_Default_GraphNodeKernel.hpp>
#include <impl/Kokkos_Default_GraphNode_Impl.hpp>

#endif  // KOKKOS_HOST_GRAPH_IMPL_HPP
