
#ifndef PLS_DATAFLOW_INTERNAL_FUNCTION_NODE_IMPL_H_
#define PLS_DATAFLOW_INTERNAL_FUNCTION_NODE_IMPL_H_

#include "graph.h"
#include "pls/internal/helpers/easy_profiler.h"

namespace pls {
namespace dataflow {
namespace internal {

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

template<typename I0, typename ...I, typename O0, typename ...O, typename F>
template<typename ...IS>
void function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
operator>>(graph<inputs<IS...>, outputs<O0, O...>> &graph) {
  out_port_ >> graph.output_ports();
}

template<typename I0, typename ...I, typename O0, typename ...O, typename F>
template<int POS, typename T>
void function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
token_pushed(token<T> token) {
  auto current_memory = get_invocation<invocation_memory>(token.invocation());

  std::get<POS>(current_memory->input_buffer_) = token;
  auto remaining_inputs = --(current_memory->inputs_missing_);
  if (remaining_inputs == 0) {
    execute_function(current_memory, token.invocation());
    current_memory->inputs_missing_ = num_in_ports;
  }
}

//////////////////////////////////////////////////////////////////
// Helpers for actually calling the work lambda
//////////////////////////////////////////////////////////////////
template<typename I0, typename ...I, typename O0, typename ...O, typename F>
void function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
execute_function(invocation_memory *invocation_memory, invocation_info invocation_info) {
  auto lambda = [=]() {
    PROFILE_WORK_BLOCK("Function Node")
    input_tuple &inputs = invocation_memory->input_buffer_;
    output_tuple outputs{};

    execute_function_internal(inputs, typename sequence_gen<1 + sizeof...(I)>::type(),
                              outputs, typename sequence_gen<1 + sizeof...(O)>::type(), invocation_info);
    PROFILE_END_BLOCK
  };
  // TODO: maybe replace this with 'continuation' style invocation
  pls::scheduler::spawn_child<pls::lambda_task_by_value<decltype(lambda)>>(lambda);
}

template<typename I0, typename ...I, typename O0, typename ...O, typename F>
template<int N, typename ...OT>
struct function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
propagate_output {
  propagate_output(multi_out_port_type &, output_tuple &, invocation_info&) {}
  void propagate() {}
};
template<typename I0, typename ...I, typename O0, typename ...O, typename F>
template<int N, typename OT1, typename ...OT>
struct function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
propagate_output<N, OT1, OT...> {
  multi_out_port_type &out_port_;
  output_tuple &output_tuple_;
  invocation_info &invocation_info_;

  propagate_output(multi_out_port_type &out_port, output_tuple &output_tuple, invocation_info&  invocation_info) :
    out_port_{out_port}, output_tuple_{output_tuple}, invocation_info_{invocation_info} {}

  void propagate() {
    std::get<N>(output_tuple_).set_invocation(invocation_info_);
    out_port_.template get<N>().push_token(std::get<N>(output_tuple_));
    propagate_output<N + 1, OT...>{out_port_, output_tuple_, invocation_info_}.propagate();
  }
};

template<typename I0, typename ...I, typename O0, typename ...O, typename F>
template<int ...IS, int ...OS>
void function_node<inputs<I0, I...>, outputs<O0, O...>, F>::
execute_function_internal(input_tuple &inputs, sequence<IS...>,
                          output_tuple &outputs, sequence<OS...>,
                          invocation_info invocation_info) {
  function_(std::get<IS>(inputs).value()..., std::get<OS>(outputs).value()...);
  propagate_output<0, O0, O...>{out_port_, outputs, invocation_info}.propagate();
}

}
}
}

#endif //PLS_DATAFLOW_INTERNAL_FUNCTION_NODE_IMPL_H_
