#ifndef PLS_SCHEDULER_MEMORY_H
#define PLS_SCHEDULER_MEMORY_H

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

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

namespace pls {
namespace internal {
namespace scheduling {

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

class scheduler_memory {
 private:
  size_t max_threads_;
  thread_state **thread_states_;
  scheduler_thread **threads_;
  data_structures::aligned_stack **task_stacks_;

 protected:
  void init(size_t max_therads,
            thread_state **thread_states,
            scheduler_thread **threads,
            data_structures::aligned_stack **task_stacks) {
    max_threads_ = max_therads;
    thread_states_ = thread_states;
    threads_ = threads;
    task_stacks_ = task_stacks;
  }

 public:
  size_t max_threads() const {
    return max_threads_;
  }
  thread_state *thread_state_for(size_t id) const {
    return thread_states_[id];
  }
  scheduler_thread *thread_for(size_t id) const {
    return threads_[id];
  }
  data_structures::aligned_stack *task_stack_for(size_t id) const {
    return task_stacks_[id];
  }
};

template<size_t MAX_THREADS, size_t TASK_STACK_SIZE>
class static_scheduler_memory : public scheduler_memory {
  // Everyone of these types has to live on its own cache line,
  // as each thread uses one of them independently.
  // Therefore it would be a major performance hit if we shared cache lines on these.
  using aligned_thread = base::alignment::aligned_wrapper<scheduler_thread>;
  using aligned_thread_state = base::alignment::aligned_wrapper<thread_state>;
  using aligned_thread_stack = base::alignment::aligned_wrapper<std::array<char, TASK_STACK_SIZE>>;
  using aligned_aligned_stack = base::alignment::aligned_wrapper<data_structures::aligned_stack>;

  // Actual Memory
  std::array<aligned_thread, MAX_THREADS> threads_;
  std::array<aligned_thread_state, MAX_THREADS> thread_states_;
  std::array<aligned_thread_stack, MAX_THREADS> task_stacks_memory_;
  std::array<aligned_aligned_stack, MAX_THREADS> task_stacks_;

  // References for parent
  std::array<scheduler_thread *, MAX_THREADS> thread_refs_;
  std::array<thread_state *, MAX_THREADS> thread_state_refs_;
  std::array<data_structures::aligned_stack *, MAX_THREADS> task_stack_refs_;

 public:
  static_scheduler_memory() : scheduler_memory() {
    for (size_t i = 0; i < MAX_THREADS; i++) {
      new((void *) task_stacks_[i].pointer()) data_structures::aligned_stack(task_stacks_memory_[i].pointer()->data(),
                                                                             TASK_STACK_SIZE);

      thread_refs_[i] = threads_[i].pointer();
      thread_state_refs_[i] = thread_states_[i].pointer();
      task_stack_refs_[i] = task_stacks_[i].pointer();
    }

    init(MAX_THREADS, thread_state_refs_.data(), thread_refs_.data(), task_stack_refs_.data());
  }
};

class malloc_scheduler_memory : public scheduler_memory {
  // Everyone of these types has to live on its own cache line,
  // as each thread uses one of them independently.
  // Therefore it would be a major performance hit if we shared cache lines on these.
  using aligned_thread = base::alignment::aligned_wrapper<scheduler_thread>;
  using aligned_thread_state = base::alignment::aligned_wrapper<thread_state>;
  using aligned_aligned_stack = base::alignment::aligned_wrapper<data_structures::aligned_stack>;

  const size_t num_threads_;

  // Actual Memory
  aligned_thread *threads_;
  aligned_thread_state *thread_states_;
  char **task_stacks_memory_;
  aligned_aligned_stack *task_stacks_;

  // References for parent
  scheduler_thread **thread_refs_;
  thread_state **thread_state_refs_;
  data_structures::aligned_stack **task_stack_refs_;

 public:
  explicit malloc_scheduler_memory(size_t num_threads, size_t memory_per_stack = 2 << 16);
  ~malloc_scheduler_memory();
};

}
}
}

#endif //PLS_SCHEDULER_MEMORY_H
