#ifndef PLS_TASK_MANAGER_H_
#define PLS_TASK_MANAGER_H_

#include <memory>
#include <utility>
#include <array>

#include "context_switcher/continuation.h"

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

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

namespace pls {
namespace internal {
namespace scheduling {

/**
 * Handles management of tasks in the system. Each thread has a local task manager,
 * responsible for allocating, freeing and publishing tasks for stealing.
 */
class task_manager {
 public:
  explicit task_manager(task *tasks,
                        data_structures::aligned_stack static_stack_space,
                        size_t num_tasks,
                        size_t stack_size) {
    for (size_t i = 0; i < num_tasks - 1; i++) {
      tasks[i].init(static_stack_space.push_bytes(stack_size), stack_size, i, 0);
      if (i > 0) {
        tasks[i].set_prev(&tasks[i - 1]);
      }
      if (i < num_tasks - 2) {
        tasks[i].set_next(&tasks[i + 1]);
      }
    }

    num_tasks_ = num_tasks;
    this_thread_tasks_ = tasks;
    active_task_ = &tasks[0];
  }

  task &get_this_thread_task(size_t depth) {
    return this_thread_tasks_[depth];
  }

  void set_thread_id(unsigned id) {
    for (size_t i = 0; i < num_tasks_; i++) {
      this_thread_tasks_[i].set_thread_id(id);
    }
  }

  task &get_active_task() {
    return *active_task_;
  }

  template<typename F>
  void spawn_child(F &&lambda) {
    // TODO: Here is some potential for optimization. We could try placing everything manually on the stack.
    active_task_->get_next()->run_as_task([lambda, this](context_switcher::continuation cont) {
      auto *last_task = active_task_;
      auto *this_task = active_task_->get_next();

      last_task->set_continuation(std::move(cont));
      active_task_ = this_task;

      // TODO: Publish last task on deque (do this properly, but this simulates the fastest possible impl)
//      last_task->flag_.store(1, std::memory_order_seq_cst);

      lambda();

      // TODO: Check if task was stolen from deque (do this properly, but this simulates the fastest possible impl)
//      if (last_task->flag_.exchange(0, std::memory_order_seq_cst) == 1) {
      active_task_ = last_task;
      return std::move(last_task->get_continuation());
//      } else {
//        return context_switcher::continuation{nullptr};
//      }
    });
  }

 private:
  size_t num_tasks_;

  task *this_thread_tasks_;
  task *active_task_;
};

template<size_t NUM_TASKS, size_t STACK_SIZE>
class static_task_manager {
 public:
  static_task_manager()
      : tasks_{},
        static_stack_storage_{},
        task_manager_{tasks_.data(), static_stack_storage_.get_stack(), NUM_TASKS, STACK_SIZE} {};
  task_manager &get_task_manager() { return task_manager_; }

 private:
  std::array<task, NUM_TASKS> tasks_;
  data_structures::static_aligned_stack<NUM_TASKS * STACK_SIZE> static_stack_storage_;

  task_manager task_manager_;
};

}
}
}
#endif //PLS_TASK_MANAGER_H_