
#ifndef PLS_DATAFLOW_INTERNAL_GRAPH_H_
#define PLS_DATAFLOW_INTERNAL_GRAPH_H_

#include <atomic>
#include <tuple>
#include <iostream>

#include "build_state.h"
#include "node.h"
#include "function_node.h"
#include "in_port.h"
#include "out_port.h"
#include "inputs.h"
#include "outputs.h"

#include "pls/pls.h"

namespace pls {
namespace dataflow {
namespace internal {

template<typename INS, typename OUTS>
class graph {};

template<typename I0, typename ...I, typename O0, typename ...O>
class graph<pls::dataflow::inputs<I0, I...>, pls::dataflow::outputs<O0, O...>> : public node {
 private:
  // Our own type
  using self_type = graph<pls::dataflow::inputs<I0, I...>, pls::dataflow::outputs<O0, O...>>;

  // Input-Output port types (internal)
  using inputs_type = multi_out_port<I0, I...>;
  using outputs_type = multi_in_port<self_type, 0, O0, O...>;

  // Input-Output value tuples
  using value_input_tuple = std::tuple<I0, I...>;
  using input_tuple = std::tuple<token < I0>, token<I>...>;

  using value_output_tuple = std::tuple<O0, O...>;
  using output_tuple = std::tuple<token < O0>, token<O>...>;

  static constexpr int num_in_ports = std::tuple_size<input_tuple>();
  static constexpr int num_out_ports = std::tuple_size<output_tuple>();

  // Memory used for ONE active invocation of the dataflow-graph
  struct invocation_memory {
    value_output_tuple *output_buffer_;
  };

 public:
  template<int POS>
  using input_at = typename inputs_type::template out_port_type_at<POS>;
  template<int POS>
  using output_at = typename outputs_type::template in_port_type_at<POS>;

  template<int POS>
  input_at<POS> &input() {
    return inputs_.template get<POS>();
  }
  template<int POS>
  output_at<POS> &output() {
    return outputs_.template get<POS>();
  }

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

  graph() : inputs_{}, outputs_{this, this} {}

  /////////////////////////////////////////////////////////////////
  // Graph building
  /////////////////////////////////////////////////////////////////
  void 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;
  }

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

  void 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<int N, typename ...IT>
  struct feed_inputs {
    feed_inputs(inputs_type &, value_input_tuple &, invocation_info &) {}
    void run() {}
  };
  template<int N, typename IT1, typename ...IT>
  struct 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();
    }
  };

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

    void execute_internal() override {
      feed_inputs<0, I0, I...>{self_->inputs_, input_, invocation_}.run();
    }
  };

  void run(value_input_tuple input, value_output_tuple &output) {
    pls::scheduler::spawn_child<run_graph_task>(this, input, &output);
  }

  //////////////////////////////////////////////////////////////////
  // Overrides for generic node functionality (building the graph)
  //////////////////////////////////////////////////////////////////
  int num_successors() const override {
    return 0;
  }
  node *successor_at(int pos) const override {
    PLS_ERROR("A graph instance has no direct successor!")
  }

  bool is_fully_connected() const override {
    return inputs_.fully_connected() && outputs_.fully_connected();
  }

  int instance_buffer_size() const override {
    return sizeof(invocation_memory);
  }
  void init_instance_buffer(void *memory) const override {
    auto invocation = new(memory) invocation_memory{};
  };

 private:
  inputs_type inputs_;
  outputs_type outputs_;

  // Information about building the graph
  node *node_list_start_{nullptr};
  node *node_list_current_{nullptr};
  int num_nodes_{0};
};

}
}
}

#endif //PLS_DATAFLOW_INTERNAL_GRAPH_H_
