
#ifndef PLS_INTERNAL_SCHEDULING_CONT_H_
#define PLS_INTERNAL_SCHEDULING_CONT_H_

#include <type_traits>
#include <atomic>
#include <utility>

#include "pls/internal/data_structures/stamped_integer.h"
#include "pls/internal/data_structures/delayed_initialization.h"
#include "pls/internal/base/alignment.h"
#include "pls/internal/base/error_handling.h"

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

#include "parallel_result.h"
#include "memory_block.h"

namespace pls {
namespace internal {
namespace scheduling {

class base_cont {
 protected:
  // We plan to only init the members for a continuation on the slow path.
  // If we can execute everything inline we simply skip it saving runtime overhead.
  template<typename T>
  using delayed_init = data_structures::delayed_initialization<T>;

 public:
  explicit base_cont(base_cont *parent, memory_block *memory_block, bool is_right_child)
      : parent_{parent},
        memory_block_{memory_block},
        is_right_child_{is_right_child} {
    PLS_ASSERT(parent_ == nullptr || parent_->memory_block_->get_depth() == memory_block_->get_depth() - 1,
               "Must only build cont chains with matching depth!")
  };

  /**
   * Execute the continuation itself.
   * Make sure to only call when all required results are in.
   * Will store the result in it's parent, but not mess with any counters.
   */
  virtual void execute() = 0;

  /**
   * Execute the right hand side task associated with the continuation.
   * Will store the result in it's parent, but not mess with any counters.
   */
  virtual void execute_task() = 0;
  virtual base_task *get_task() = 0;

  virtual void *get_right_result_pointer() = 0;
  virtual void *get_left_result_pointer() = 0;

  template<typename T>
  void store_right_result(T &&result) {
    using BASE_T = typename std::remove_cv<typename std::remove_reference<T>::type>::type;
    reinterpret_cast<delayed_init<BASE_T> *>(get_right_result_pointer())->initialize(std::forward<T>(result));
  }

  template<typename T>
  void store_left_result(T &&result) {
    using BASE_T = typename std::remove_cv<typename std::remove_reference<T>::type>::type;
    reinterpret_cast<delayed_init<BASE_T> *>(get_left_result_pointer())->initialize(std::forward<T>(result));
  }

  base_cont *get_parent() { return parent_; }
  memory_block *get_memory_block() { return memory_block_; }
  bool is_right_child() const { return is_right_child_; }

 protected:
  base_cont *parent_;
  memory_block *memory_block_;
  bool is_right_child_;
};

template<typename T2, typename R1, typename R2, typename F>
class cont : public base_cont {
 private:
  template<typename RES_TYPE>
  struct result_runner {
    // Strip off unwanted modifiers...
    using BASE_RES_TYPE = typename std::remove_cv<typename std::remove_reference<RES_TYPE>::type>::type;

    static void execute(cont &cont) {
      parallel_result<BASE_RES_TYPE>
          result{cont.function_((*cont.left_result_).value(), (*cont.right_result_).value())};
      if (result.fast_path() && cont.parent_ != nullptr) {
        if (cont.is_right_child()) {
          cont.parent_->store_right_result(std::move(result));
        } else {
          cont.parent_->store_left_result(std::move(result));
        }
      }
    }
  };
  template<typename INNER_TYPE>
  struct result_runner<parallel_result<INNER_TYPE>> {
    static void execute(cont &cont) {
      auto result = cont.function_((*cont.left_result_).value(), (*cont.right_result_).value());
      if (result.fast_path() && cont.parent_) {
        if (cont.is_right_child()) {
          cont.parent_->store_right_result(std::move(result));
        } else {
          cont.parent_->store_left_result(std::move(result));
        }
      }
    }
  };

 public:
  template<typename FARG, typename ...T2ARGS>
  explicit cont(base_cont *parent,
                memory_block *memory_block,
                bool is_right_child,
                FARG &&function,
                T2ARGS...task_2_args):
      base_cont(parent, memory_block, is_right_child),
      function_{std::forward<FARG>(function)},
      task_{std::forward<T2ARGS>(task_2_args)..., this} {};

  void execute() override {
    PROFILE_CONTINUATION("execute_cont");
    using result_type = decltype(function_((*left_result_).value(), (*right_result_).value()));
    result_runner<result_type>::execute(*this);

    this->~cont();
    auto *memory_block = this->get_memory_block();
    memory_block->free_buffer();
    memory_block->reset_state();
  }

  void execute_task() override {
    task_.execute();
  }
  base_task *get_task() override {
    return &task_;
  }

  void *get_left_result_pointer() override {
    return &left_result_;
  }
  void *get_right_result_pointer() override {
    return &right_result_;
  }

 private:
  // Initial data members. These slow down the fast path, try to init them lazy when possible.
  F function_;
  T2 task_;

  // Some fields/actual values stay uninitialized (save time on the fast path if we don not need them).
  // More fields untouched on the fast path is good, but for ease of an implementation we only keep some for now.
  delayed_init<R1> left_result_;
  delayed_init<R2> right_result_;
};

}
}
}

#endif //PLS_INTERNAL_SCHEDULING_CONT_H_
