
#ifndef PLS_CONT_MANAGER_H_
#define PLS_CONT_MANAGER_H_

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

#include "pls/internal/data_structures/aligned_stack.h"
#include "pls/internal/scheduling/cont.h"
#include "pls/internal/scheduling/thread_state.h"

namespace pls {
namespace internal {
namespace scheduling {

class cont_manager {
 public:
  // Helper to pass the compile time constants to the constructor.
  template<size_t NUM_CONTS, size_t MAX_CONT_SIZE>
  struct template_args {};

  template<size_t NUM_CONTS, size_t MAX_CONT_SIZE>
  explicit cont_manager(data_structures::aligned_stack &cont_storage, template_args<NUM_CONTS, MAX_CONT_SIZE>)
      : max_cont_size_{MAX_CONT_SIZE}, num_conts_{NUM_CONTS} {

    // First node is currently active and our local start
    active_node_ = init_memory_block<MAX_CONT_SIZE>(cont_storage, nullptr, 0);

    // Build up chain after it
    memory_block *current_node = active_node_;
    for (size_t i = 1; i < NUM_CONTS; i++) {
      memory_block *next_node = init_memory_block<MAX_CONT_SIZE>(cont_storage, current_node, i);
      current_node->set_next(next_node);
      current_node = next_node;
    }
  };

  // Aquire and release memory blocks...
  memory_block *get_next_memory_block() {
    auto result = active_node_;
    active_node_ = active_node_->get_next();
    return result;
  }
  void return_memory_block() {
    active_node_ = active_node_->get_prev();
  }
  void move_active_node(int depth) {
    if (depth < 0) {
      for (long i = 0; i < (depth * -1); i++) {
        active_node_ = active_node_->get_prev();
      }
    } else {
      for (long i = 0; i < depth; i++) {
        active_node_ = active_node_->get_next();
      }
    }
  }
  void move_active_node_to_start() {
    move_active_node(-1 * active_node_->get_depth());
  }
  memory_block *get_active_node() {
    return active_node_;
  }

  bool is_clean() {
    if (get_active_node()->get_depth() == 0) {
      memory_block *current_node = active_node_;
      for (size_t i = 1; i < num_conts_; i++) {
        if (current_node->get_prev() != nullptr && current_node->get_prev()->get_next() != current_node) {
          return false;
        }
        if (current_node->is_buffer_used()) {
          return false;
        }
        current_node = current_node->get_next();
      }
    } else {
      return false;
    }

    return true;
  }

  // Manage the fall through behaviour/slow path behaviour
  bool falling_through() const {
    return fall_through_;
  }
  void fall_through_and_notify_cont(base_cont *notified_cont, bool is_child_right) {
    fall_through_ = true;
    fall_through_cont_ = notified_cont;
    fall_through_child_right = is_child_right;
  }

  void aquire_memory_chain(memory_block *target_chain) {
    PLS_ASSERT(active_node_->get_depth() == target_chain->get_depth() + 1,
               "Can only steal aquire chain parts with correct depth.");

    active_node_->set_prev(target_chain);
    target_chain->set_next(active_node_);
  }

  void execute_fall_through_code() {
    PLS_ASSERT(falling_through(), "Must be falling through to execute the associated code.")

    auto &my_state = thread_state::get();

    // Copy fall through status and reset it (for potentially nested execution paths).
    auto *notified_cont = fall_through_cont_;

    fall_through_cont_ = nullptr;
    fall_through_ = false;

    // Keep the target chain before we execute, as this potentially frees the memory
    auto *target_memory_block = notified_cont->get_memory_block();
    auto *target_chain = target_memory_block->get_offered_chain().load();

    // Notify the next continuation of finishing a child...
    if (target_memory_block->get_results_missing().fetch_add(-1) == 1) {
      // ... we finished the continuation.
      // We are now in charge continuing to execute the above continuation chain.

      PLS_ASSERT(active_node_->get_prev()->get_depth() == target_memory_block->get_depth(),
                 "We must hold the system invariant to be in the correct depth.")
      if (active_node_->get_prev() != target_memory_block) {
        // We do not own the thing we will execute.
        // Own it by swapping the chain belonging to it in.
        aquire_memory_chain(target_memory_block);
      }
      my_state.parent_cont_ = notified_cont->get_parent();
      my_state.right_spawn_ = notified_cont->is_right_child();
      active_node_ = target_memory_block;
      notified_cont->execute();
      if (!falling_through() && notified_cont->get_parent() != nullptr) {
        fall_through_and_notify_cont(notified_cont->get_parent(), notified_cont->is_right_child());
      }
      return;
    } else {
      // ... we did not finish the last continuation.
      // We are no longer in charge of executing the above continuation chain.

      PLS_ASSERT(active_node_->get_prev()->get_depth() == target_memory_block->get_depth(),
                 "We must hold the system invariant to be in the correct depth.")
      if (active_node_->get_prev() == target_memory_block) {
        // We own the thing we are not allowed to execute.
        // Get rid of the ownership by using the offered chain.
        aquire_memory_chain(target_chain);
      }

      move_active_node_to_start();
      // We are done here...nothing more to execute
      return;
    }
  }

 private:
  template<size_t MAX_CONT_SIZE>
  static memory_block *init_memory_block(data_structures::aligned_stack &cont_storage,
                                         memory_block *prev,
                                         unsigned long depth) {
    // Represents one cont_node and its corresponding memory buffer (as one continuous block of memory).
    constexpr size_t buffer_size = MAX_CONT_SIZE - base::alignment::next_alignment(sizeof(memory_block));
    char *memory_block_ptr = cont_storage.push_bytes<memory_block>();
    char *memory_block_buffer_ptr = cont_storage.push_bytes(buffer_size);

    return new(memory_block_ptr) memory_block(memory_block_buffer_ptr, buffer_size, prev, depth);
  }

 private:
  const size_t max_cont_size_;
  const size_t num_conts_;

  /**
   * Managing the continuation chain.
   */
  memory_block *active_node_;

  /**
   * Managing falling through back to the scheduler.
   */
  bool fall_through_{false};
  bool fall_through_child_right{false};
  base_cont *fall_through_cont_{nullptr};
};

template<size_t NUM_CONTS, size_t MAX_CONT_SIZE>
class static_cont_manager {
 public:
  static_cont_manager()
      : static_cont_storage_{},
        cont_manager_(static_cont_storage_.get_stack(), cont_manager::template_args<NUM_CONTS, MAX_CONT_SIZE>{}) {}
  cont_manager &get_cont_manager() { return cont_manager_; }

 private:
  data_structures::static_aligned_stack<NUM_CONTS * MAX_CONT_SIZE> static_cont_storage_;
  cont_manager cont_manager_;
};

}
}
}

#endif //PLS_CONT_MANAGER_H_
