#ifndef PLS_SCHEDULER_H
#define PLS_SCHEDULER_H

#include <array>
#include <iostream>

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

#include "pls/internal/data_structures/aligned_stack.h"

#include "pls/internal/base/thread.h"
#include "pls/internal/base/barrier.h"

#include "pls/internal/scheduling/thread_state.h"
#include "pls/internal/scheduling/scheduler_memory.h"

namespace pls {
namespace internal {
namespace scheduling {

void worker_routine();
using scheduler_thread = base::thread<decltype(&worker_routine), thread_state>;

class scheduler {
  friend class task;
  friend void worker_routine();

  const unsigned int num_threads_;
  scheduler_memory *memory_;

  base::barrier sync_barrier_;
  bool terminated_;
 public:
  /**
   * Initializes a scheduler instance with the given number of threads.
   * This will spawn the threads and put them to sleep, ready to process an
   * upcoming parallel section.
   *
   * @param memory All memory is allocated statically, thus the user is required to provide the memory instance.
   * @param num_threads The number of worker threads to be created.
   */
  explicit scheduler(scheduler_memory *memory, unsigned int num_threads);

  /**
   * The scheduler is implicitly terminated as soon as it leaves the scope.
   */
  ~scheduler();

  /**
   * Wakes up the thread pool.
   * Code inside the Function lambda can invoke all parallel APIs.
   * This is meant to cleanly sleep and wake up the scheduler during an application run,
   * e.g. to run parallel code on a timer loop/after interrupts.
   *
   * @param work_section generic function or lambda to be executed in the scheduler's context.
   */
  template<typename Function>
  void perform_work(Function work_section);

  /**
   * Explicitly terminate the worker threads. Scheduler must not be used after this.
   *
   * @param wait_for_workers Set to true if you wish to return from this method only after the workers are shut down.
   */
  void terminate(bool wait_for_workers = true);

  /**
   * Helper to spawn a child on the currently running task.
   *
   * @tparam T type of the new task
   * @param sub_task the new task to be spawned
   */
  template<typename T>
  void spawn_child(T &sub_task) {
    this_thread_state()->current_task_->spawn_child(sub_task);
  }

  /**
   * Helper to wait for all children of the currently executing task.
   */
  void wait_for_all() {
    this_thread_state()->current_task_->wait_for_all();
  }

  unsigned int num_threads() const { return num_threads_; }

 private:
  // Helpers for accessing thread states
  thread_state *thread_state_for(size_t id) { return memory_->thread_state_for(id); }

  task *get_local_task();
  task *steal_task();
};

}
}
}
#include "scheduler_impl.h"

#endif //PLS_SCHEDULER_H