
#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
    start_node_ = active_node_ = init_memory_block<MAX_CONT_SIZE>(cont_storage, nullptr, nullptr, 0);

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

  // Aquire and release memory blocks...
  memory_block *get_next_memory_block() {
    active_node_->set_start(start_node_);
    auto result = active_node_;

    active_node_ = active_node_->get_next();

    return result;
  }
  void return_memory_block() {
    active_node_ = active_node_->get_prev();
  }

  // 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) {
    auto *our_next_node = get_node(target_chain->get_depth() + 1);

    our_next_node->set_prev(target_chain);
    target_chain->set_next(our_next_node);

    start_node_ = target_chain->get_start();
    active_node_ = target_chain;
  }

  memory_block *get_node(unsigned int depth) {
    // TODO: Remove this O(n) factor to avoid the
    //       T_1/P + (T_lim + D) stealing time bound.
    memory_block *current = start_node_;
    for (unsigned int i = 0; i < depth; i++) {
      current->set_start(start_node_);
      current = current->get_next();
    }

    return current;
  }

  void check_clean_chain() {
    memory_block *current = start_node_;
    for (unsigned int i = 0; i < num_conts_; i++) {
      bool buffer_used = current->is_buffer_used();
      auto state_value = current->get_state().load().value;
      if (state_value != memory_block::initialized || buffer_used || current->get_depth() != i) {
        PLS_ASSERT(false,
                   "Must always steal with a clean chain!");
      }

      current->set_start(start_node_);
      current = current->get_next();
    }
  }

  void set_active_depth(unsigned int depth) {
    active_node_ = get_node(depth);
  }

  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_;
    bool notifier_is_right_child = fall_through_child_right;

    std::cout << "Notifying Cont on core " << my_state.get_id() << " and depth "
              << notified_cont->get_memory_block()->get_depth() << std::endl;

    fall_through_cont_ = nullptr;
    fall_through_ = false;

    // Special case for lock free implementation.
    // When we finish a 'left strain' it can happen that the 'right strain'
    // is currently being stolen. We need to be sure that this steal is finished
    // as we need the thieves memory blocks in case we get blocked by it.
    if (!notifier_is_right_child) {
      // Check to execute right child directly...
      auto &atomic_state = notified_cont->get_memory_block()->get_state();
      auto old_state = atomic_state.load();
      memory_block::stamped_state target_state{old_state.stamp + 1, memory_block::state::execute_local};
      memory_block::stamped_state exchanged_state = atomic_state.exchange(target_state);
      if (exchanged_state.value != memory_block::state::stolen) {
        // We 'overruled' the stealing thread and execute the other task ourselfs.
        // We can be sure that only we where involved in executing the child tasks of the parent_continuation...
        notified_cont->get_memory_block()->get_results_missing().fetch_add(-1);

        my_state.parent_cont_ = notified_cont;
        my_state.right_spawn_ = true;
        notified_cont->execute_task();
        if (falling_through()) {
          // ... if the second strain was interrupted we fall through without scheduling the parent_continuation
          // (the currently pending/interrupted strain will do this itself).
          return;
        } else {
          // ... we could finish the second strain.
          // Register the parent continuation for being notified.
          // (This is not the most efficient, as we could simply execute it. However,
          // this way of doing it spares us from duplicating a lot of code).
          fall_through_and_notify_cont(notified_cont, true);
          return;
        }
      }

      // Right side is 'fully' stolen. We can continue to inform the parent like we would do normally.
    }

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

    // Notify the next continuation of finishing a child...
    if (notified_cont->get_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.

      if (get_node(notified_cont->get_memory_block()->get_depth()) != notified_cont->get_memory_block()) {
        my_state.cont_manager_.check_clean_chain();
        // We do not own the thing we will execute.
        // Own it by swapping the chain belonging to it in.
        aquire_memory_chain(notified_cont->get_memory_block());
        std::cout << "Now in charge of memory chain on core " << my_state.get_id() << std::endl;
      }
      my_state.parent_cont_ = notified_cont->get_parent();
      my_state.right_spawn_ = notified_cont->is_right_child();
      active_node_ = notified_cont->get_memory_block();
      std::cout << "Execute cont on core " << my_state.get_id() << std::endl;
      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.

      if (get_node(notified_cont->get_memory_block()->get_depth()) == notified_cont->get_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);
        std::cout << "No longer in charge of chain above on core " << my_state.get_id() << std::endl;
      }

      my_state.cont_manager_.check_clean_chain();
      // 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 *memory_chain_start,
                                         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, memory_chain_start, prev, depth);
  }

 private:
  const size_t max_cont_size_;
  const size_t num_conts_;

  /**
   * Managing the continuation chain.
   */
  memory_block *start_node_;
  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_
