
#ifndef PLS_DATAFLOW_INTERNAL_GRAPH_IMPL_H_
#define PLS_DATAFLOW_INTERNAL_GRAPH_IMPL_H_

#include "pls/internal/helpers/profiler.h"

namespace pls {
namespace dataflow {
namespace internal {

template<typename I0, typename ...I, typename O0, typename ...O>
template<typename ...OS, typename FUNC>
function_node<inputs<I0, I...>, outputs<OS...>, FUNC> &graph<inputs<I0, I...>, outputs<O0, O...>>::
operator>>(function_node<inputs<I0, I...>, outputs<OS...>, FUNC> &other_node) {
  inputs_ >> other_node.in_ports();
  return other_node;
}

template<typename I0, typename ...I, typename O0, typename ...O>
template<int POS, typename T>
void graph<inputs<I0, I...>, outputs<O0, O...>>::
token_pushed(token<T> token) {
  auto invocation = get_invocation<invocation_memory>(token.invocation());
  std::get<POS>(*invocation->output_buffer_) = token.value();
}

template<typename I0, typename ...I, typename O0, typename ...O>
void graph<inputs<I0, I...>, outputs<O0, O...>>::
build() {
  PLS_ASSERT(build_state_ == build_state::fresh, "Must only build a dataflow graph once!")
  PLS_ASSERT(is_fully_connected(), "Must fully connect all inputs/outputs inside a dataflow graph!")

  build_state_ = build_state::building;
  node_list_start_ = node_list_current_ = this;
  memory_index_ = 0;
  num_nodes_ = 1;
  for (int i = 0; i < num_in_ports; i++) {
    build_recursive(inputs_.next_node_at(i));
  }
  build_state_ = build_state::built;
}

template<typename I0, typename ...I, typename O0, typename ...O>
void graph<inputs<I0, I...>, outputs<O0, O...>>::
build_recursive(node *node) {
  if (node->build_state_ != build_state::fresh) {
    return; // Already visited
  }
  node->build_state_ = build_state::building;
  PLS_ASSERT(node->is_fully_connected(), "Must fully connect dataflow graph nodes!")

  add_node(node);
  for (int i = 0; i < node->num_successors(); i++) {
    build_recursive(node->successor_at(i));
  }

  node->build_state_ = build_state::built;
}

template<typename I0, typename ...I, typename O0, typename ...O>
void graph<inputs<I0, I...>, outputs<O0, O...>>::
add_node(node *new_node) {
  new_node->memory_index_ = num_nodes_++;

  if (node_list_current_ == nullptr) {
    node_list_current_ = new_node;
    node_list_start_ = new_node;
  } else {
    node_list_current_->direct_successor_ = new_node;
    node_list_current_ = new_node;
  }
}

template<typename I0, typename ...I, typename O0, typename ...O>
template<int N, typename ...IT>
struct graph<inputs<I0, I...>, outputs<O0, O...>>::
feed_inputs {
  feed_inputs(inputs_type &, value_input_tuple &, invocation_info &) {}
  void run() {}
};

template<typename I0, typename ...I, typename O0, typename ...O>
template<int N, typename IT1, typename ...IT>
struct graph<inputs<I0, I...>, outputs<O0, O...>>::
feed_inputs<N, IT1, IT...> {
  inputs_type &inputs_;
  value_input_tuple &input_values_;
  invocation_info &invocation_;

  feed_inputs(inputs_type &inputs,
              value_input_tuple &input_values,
              invocation_info &invocation) : inputs_{inputs},
                                             input_values_{input_values},
                                             invocation_{invocation} {}

  void run() {
    inputs_.template get<N>().push_token(token<IT1>{std::get<N>(input_values_), invocation_});
    feed_inputs<N + 1, IT...>{inputs_, input_values_, invocation_}.run();
  }
};

template<typename I0, typename ...I, typename O0, typename ...O>
class graph<inputs<I0, I...>, outputs<O0, O...>>::run_graph_task : public pls::task {
  graph *self_;
  value_input_tuple input_;
  value_output_tuple *output_;

  // Buffers for actual execution
  invocation_info invocation_;

 public:
  run_graph_task(graph *self, value_input_tuple input, value_output_tuple *output) : self_{self},
                                                                                      input_{input},
                                                                                      output_{output},
                                                                                      invocation_{nullptr} {
    void **buffers;
    buffers = reinterpret_cast<void **>(allocate_memory(self_->num_nodes_ * sizeof(void *)));

    node *iterator = self_->node_list_start_;
    for (int i = 0; i < self_->num_nodes_; i++) {
      auto required_size = iterator->instance_buffer_size();
      buffers[i] = allocate_memory(required_size);
      iterator->init_instance_buffer(buffers[i]);

      iterator = iterator->direct_successor_;
    }
    invocation_ = invocation_info{buffers};

    self_->get_invocation<invocation_memory>(invocation_)->output_buffer_ = output;
  }

  ~run_graph_task() {
    node *iterator = self_->node_list_start_;
    for (int i = 0; i < self_->num_nodes_; i++) {
      void* memory = invocation_.get_instance_buffer<void>(i);
      iterator->clean_up_instance_buffer(memory);

      iterator = iterator->direct_successor_;
    }
  }

  void execute_internal() override {
    PROFILE_WORK_BLOCK("Graph Invocation")
    feed_inputs<0, I0, I...>{self_->inputs_, input_, invocation_}.run();
    
    wait_for_all();
    this->~run_graph_task();
    PROFILE_END_BLOCK
  }
};

}
}
}

#endif //PLS_DATAFLOW_INTERNAL_GRAPH_IMPL_H_
