
#ifndef PLS_SCHEDULER_IMPL_H
#define PLS_SCHEDULER_IMPL_H

#include <utility>
#include "pls/internal/scheduling/cont.h"
#include "pls/internal/scheduling/parallel_result.h"
#include "pls/internal/scheduling/task.h"

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

namespace pls {
namespace internal {
namespace scheduling {

template<typename F1, typename F2>
struct scheduler::starter {
  F1 function_1_;
  F2 function_2_;
  using return_type_1 = decltype(function_1_());
  using return_type_2 = decltype(function_2_());

  // Enforce correct return types of lambdas (parallel_result)
  static_assert(std::is_base_of<parallel_result_base, return_type_1>::value,
                "Must only return parallel results in parallel code");
  static_assert(std::is_base_of<parallel_result_base, return_type_2>::value,
                "Must only return parallel results in parallel code");

  template<typename F1ARG, typename F2ARG>
  explicit starter(F1ARG &&function_1, F2ARG &&function_2) : function_1_{std::forward<F1ARG>(function_1)},
                                                             function_2_{std::forward<F2ARG>(function_2)} {};

  template<typename FCONT>
  auto then(FCONT &&cont_function)
  -> decltype(cont_function(std::declval<typename return_type_1::value_type>(),
                            std::declval<typename return_type_2::value_type>())) {
    PROFILE_FAST_PATH("then");
    using continuation_type = cont<task<F2>, return_type_1, return_type_2, FCONT>;
    using result_type = decltype(cont_function(std::declval<typename return_type_1::value_type>(),
                                               std::declval<typename return_type_2::value_type>()));

    auto &my_state = thread_state::get();
    auto &cont_manager = my_state.get_cont_manager();

    // Select current memory block.
    // For now directly copy both the continuation function and the second task.
    // (We might optimize this in the future to require less memory copies.)
    auto *current_memory_block = cont_manager.get_next_memory_block();


    // We keep track of the last spawn to build up the parent_cont chain
    const bool is_right_cont = my_state.right_spawn_;
    base_cont *parent_cont = my_state.parent_cont_;

    continuation_type *current_cont = current_memory_block->place_in_buffer<continuation_type>(parent_cont,
                                                                                               current_memory_block,
                                                                                               is_right_cont,
                                                                                               cont_function,
                                                                                               function_2_);
    my_state.parent_cont_ = current_cont;

    // Publish the second task.
    my_state.get_task_manager().publish_task(current_cont->get_task());

    // Call first function on fast path
    my_state.right_spawn_ = false;
    return_type_1 result_1 = function_1_();
    if (!result_1.fast_path()) {
      // Get our replacement from the task stack and store it for later use when we are actually blocked.
      auto traded_memory = my_state.get_task_manager().try_pop_local();
      current_cont->get_memory_block()->get_offered_chain().store(*traded_memory);

      // Unwind stack...
      return result_type{};
    }

    // Try to call second function on fast path
    auto traded_memory = my_state.get_task_manager().try_pop_local();
    if (traded_memory) {
      // The task got stolen...get_memory_block
      // ...but we got a memory block that can be used if we block on this one.
      current_cont->get_memory_block()->get_offered_chain().store(*traded_memory);

      // Main scheduling loop is responsible for entering the result to the slow path...
      current_cont->store_left_result(std::move(result_1));
      cont_manager.fall_through_and_notify_cont(current_cont, false);
      // Unwind stack...
      return result_type{};
    } else {
      my_state.right_spawn_ = true;
      return_type_2 result_2 = function_2_();
      if (!result_2.fast_path()) {
        // Main scheduling loop is responsible for entering the result to the slow path...
        current_cont->store_left_result(std::move(result_1));
        current_cont->get_memory_block()->get_results_missing().fetch_add(-1);
        // Unwind stack...
        return result_type{};
      }

      // We fully got all results, inline as good as possible.
      // This is the common case, branch prediction should be rather good here.

      // Just return the cont object unused and directly call the function.
      current_cont->~continuation_type();
      current_memory_block->free_buffer();
      cont_manager.return_memory_block();

      // The continuation has the same execution environment as we had for the children.
      // We need this to allow spawns in there.
      my_state.parent_cont_ = parent_cont;
      my_state.right_spawn_ = is_right_cont;

      auto cont_result = cont_function(result_1.value(), result_2.value());
      if (!cont_result.fast_path()) {
        // Unwind stack...
        return result_type{};
      }
      return cont_result;
    }
  };
};

template<typename F1, typename F2>
scheduler::starter<F1, F2> scheduler::invoke_parallel(F1 &&function_1, F2 &&function_2) {
  return scheduler::starter<F1, F2>{std::forward<F1>(function_1), std::forward<F2>(function_2)};
}

template<typename F1, typename F2>
scheduler::starter<F1, F2> scheduler::par(F1 &&function_1, F2 &&function_2) {
  return thread_state::get().get_scheduler().invoke_parallel(std::forward<F1>(function_1),
                                                             std::forward<F2>(function_2));
}

class scheduler::init_function {
 public:
  virtual void run() = 0;
};
template<typename F>
class scheduler::init_function_impl : public init_function {
 public:
  explicit init_function_impl(F &function) : function_{function} {}
  void run() override {
    scheduler::par([]() {
      return parallel_result<int>{0};
    }, [=]() {
      return function_();
    }).then([=](int, int b) {
      thread_state::get().get_scheduler().work_section_done_ = true;
      return parallel_result<int>{b};
    });

  }
 private:
  F &function_;
};

template<typename Function>
void scheduler::perform_work(Function work_section) {
  // Prepare main root task
  init_function_impl<Function> starter_function{work_section};
  main_thread_starter_function_ = &starter_function;

  work_section_done_ = false;
  if (reuse_thread_) {
    auto &my_state = memory_.thread_state_for(0);
    base::this_thread::set_state(&my_state); // Make THIS THREAD become the main worker

    sync_barrier_.wait(); // Trigger threads to wake up
    work_thread_work_section(); // Simply also perform the work section on the main loop
    sync_barrier_.wait(); // Wait for threads to finish
  } else {
    // Simply trigger the others to do the work, this thread will sleep/wait for the time being
    sync_barrier_.wait(); // Trigger threads to wake up
    sync_barrier_.wait(); // Wait for threads to finish
  }
}

}
}
}

#endif //PLS_SCHEDULER_IMPL_H
