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

#include "thread_state.h"

#ifndef PLS_SCHEDULER_MEMORY_H
#define PLS_SCHEDULER_MEMORY_H

namespace pls {
namespace internal {
namespace scheduling {

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

class scheduler_memory {
 public:
  virtual size_t max_threads() const = 0;
  virtual thread_state *thread_state_for(size_t id) = 0;
  virtual scheduler_thread *thread_for(size_t id) = 0;
  virtual data_structures::aligned_stack *task_stack_for(size_t id) = 0;
};

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>;

  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_;

 public:
  static_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);
    }
  }

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

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_;

  aligned_thread *threads_;
  aligned_thread_state *thread_states_;
  char **task_stacks_memory_;
  aligned_aligned_stack *task_stacks_;
 public:
  explicit malloc_scheduler_memory(size_t num_threads, size_t memory_per_stack = 2 << 16);
  ~malloc_scheduler_memory();

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

}
}
}

#endif //PLS_SCHEDULER_MEMORY_H
