
#ifndef PLS_SCHEDULER_IMPL_H
#define PLS_SCHEDULER_IMPL_H

#include <utility>

#include "context_switcher/context_switcher.h"
#include "context_switcher/continuation.h"

#include "pls/internal/scheduling/task_manager.h"
#include "pls/internal/scheduling/task.h"

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

namespace pls::internal::scheduling {

template<typename ALLOC>
scheduler::scheduler(unsigned int num_threads,
                     size_t computation_depth,
                     size_t stack_size,
                     bool reuse_thread,
                     ALLOC &&stack_allocator) :
    num_threads_{num_threads},
    reuse_thread_{reuse_thread},
    sync_barrier_{num_threads + 1 - reuse_thread},
    worker_threads_{},
    thread_states_{},
    main_thread_starter_function_{nullptr},
    work_section_done_{false},
    terminated_{false},
    stack_allocator_{std::make_shared<ALLOC>(std::forward<ALLOC>(stack_allocator))} {

  worker_threads_.reserve(num_threads);
  task_managers_.reserve(num_threads);
  thread_states_.reserve(num_threads);
  for (unsigned int i = 0; i < num_threads_; i++) {
    auto &this_task_manager =
        task_managers_.emplace_back(std::make_unique<task_manager>(i,
                                                                   computation_depth,
                                                                   stack_size,
                                                                   stack_allocator_));
    auto &this_thread_state = thread_states_.emplace_back(std::make_unique<thread_state>(*this, i, *this_task_manager));

    if (reuse_thread && i == 0) {
      worker_threads_.emplace_back();
      continue; // Skip over first/main thread when re-using the users thread, as this one will replace the first one.
    }

    auto *this_thread_state_pointer = this_thread_state.get();
    worker_threads_.emplace_back([this_thread_state_pointer] {
      thread_state::set(this_thread_state_pointer);
      work_thread_main_loop();
    });
  }
}

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 {
    auto &root_task = thread_state::get().get_task_manager().get_active_task();
    root_task.run_as_task([&](context_switcher::continuation cont) {
      thread_state::get().main_continuation() = std::move(cont);
      function_();
      thread_state::get().get_scheduler().work_section_done_.store(true);
      PLS_ASSERT(thread_state::get().main_continuation().valid(), "Must return valid continuation from main task.");
      return std::move(thread_state::get().main_continuation());
    });

  }
 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 = thread_state_for(0);
    thread_state::set(&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
  }
}

template<typename Function>
void scheduler::spawn(Function &&lambda) {
  thread_state::get().get_task_manager().spawn_child(std::forward<Function>(lambda));
}

}

#endif //PLS_SCHEDULER_IMPL_H
