
#ifndef PLS_THREAD_STATE_H
#define PLS_THREAD_STATE_H

#include <atomic>
#include <random>
#include <chrono>
#include <utility>

#include "context_switcher/continuation.h"

#include "pls/internal/build_flavour.h"

#include "pls/internal/base/system_details.h"
#include "pls/internal/data_structures/stamped_integer.h"
#include "pls/internal/base/stack_allocator.h"

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

namespace pls::internal::scheduling {

enum EMPTY_QUEUE_STATE : int32_t {
  QUEUE_NON_EMPTY = 0,
  QUEUE_MAYBE_EMPTY = 1,
  QUEUE_EMPTY = 2,
};

class scheduler;
/**
 * Proxy-Object for thread local state needed during scheduling.
 * The main use is to perform thread_state::get() as a thread local
 * memory to identify the current worker thread state.
 *
 * Holds only minimal data by itself and points to the appropriate scheduler
 * and task manager objects associated with this thread.
 */
struct PLS_CACHE_ALIGN thread_state {
 private:
  const unsigned thread_id_;
  scheduler &scheduler_;
  task_manager &task_manager_;

  context_switcher::continuation main_loop_continuation_;
  std::minstd_rand random_;
  base_task *active_task_;

  static thread_local bool is_scheduler_active_;

  std::shared_ptr<base::stack_allocator> stack_allocator_;
  char *serial_call_stack_;
  size_t serial_call_stack_size_;

#if PLS_SLEEP_WORKERS_ON_EMPTY
  PLS_CACHE_ALIGN std::atomic<data_structures::stamped_integer> queue_empty_{EMPTY_QUEUE_STATE::QUEUE_NON_EMPTY};
#endif

 public:
  explicit thread_state(scheduler &scheduler,
                        unsigned thread_id,
                        task_manager &task_manager,
                        std::shared_ptr<base::stack_allocator> &stack_allocator,
                        size_t serial_call_stack_size) :
      thread_id_{thread_id},
      scheduler_{scheduler},
      task_manager_{task_manager},
      random_{static_cast<unsigned long>(std::chrono::steady_clock::now().time_since_epoch().count()) + thread_id},
      active_task_{task_manager.get_task(0)},
      stack_allocator_{stack_allocator},
      serial_call_stack_size_{serial_call_stack_size} {
    serial_call_stack_ = stack_allocator->allocate_stack(serial_call_stack_size_);
    for (size_t i = 0; i < serial_call_stack_size; i += base::system_details::CACHE_LINE_SIZE) {
      serial_call_stack_[i] = 'a'; // Touch the stack
    }
  };

  ~thread_state() {
    stack_allocator_->free_stack(serial_call_stack_size_, serial_call_stack_);
  }

  // Do not allow accidental copy/move operations.
  thread_state(const thread_state &) = delete;
  thread_state(thread_state &&) = delete;
  thread_state &operator=(const thread_state &) = delete;
  thread_state &operator=(thread_state &&) = delete;

  /**
   * Convenience helper to get the thread_state instance associated with this thread.
   * Must only be called on threads that are associated with a thread_state,
   * this will most likely be threads created by the scheduler.
   *
   * Each call is guaranteed to be a new lookup, i.e. it is not cached after fiber context switches.
   *
   * @return The thread_state of this thread.
   */
  [[nodiscard]] static thread_state &PLS_NOINLINE get();
  static void set(thread_state *);

  [[nodiscard]] static bool is_scheduler_active() {
    return is_scheduler_active_;
  }
  static void set_scheduler_active(bool active) {
    is_scheduler_active_ = active;
  }

  [[nodiscard]] unsigned get_thread_id() const { return thread_id_; }
  [[nodiscard]] task_manager &get_task_manager() { return task_manager_; }
  [[nodiscard]] scheduler &get_scheduler() { return scheduler_; }
  [[nodiscard]] long get_rand() {
    return random_();
  }

  [[nodiscard]] context_switcher::continuation &main_continuation() {
    return main_loop_continuation_;
  }

  void set_active_task(base_task *active_task) { active_task_ = active_task; }
  [[nodiscard]] base_task *get_active_task() const { return active_task_; }

  char *get_serial_call_stack() { return serial_call_stack_; }
  [[nodiscard]] size_t get_serial_call_stack_size() const { return serial_call_stack_size_; }

#if PLS_SLEEP_WORKERS_ON_EMPTY
  [[nodiscard]] std::atomic<data_structures::stamped_integer> &get_queue_empty_flag() {
    return queue_empty_;
  };
#endif

};

}

#endif //PLS_THREAD_STATE_H
