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

#include "context_switcher/context_switcher.h"

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

#include "pls/internal/base/thread.h"
#include "pls/internal/base/error_handling.h"

namespace pls {
namespace internal {
namespace scheduling {

scheduler::scheduler(scheduler_memory &memory, const unsigned int num_threads, bool reuse_thread) :
    num_threads_{num_threads},
    reuse_thread_{reuse_thread},
    memory_{memory},
    sync_barrier_{num_threads + 1 - reuse_thread},
    terminated_{false} {
  if (num_threads_ > memory.max_threads()) {
    PLS_ERROR("Tried to create scheduler with more OS threads than pre-allocated memory.");
  }

  for (unsigned int i = 0; i < num_threads_; i++) {
    // Placement new is required, as the memory of `memory_` is not required to be initialized.
    memory.thread_state_for(i).set_scheduler(this);
    memory.thread_state_for(i).set_id(i);
    memory.thread_state_for(i).get_task_manager().set_thread_id(i);

    if (reuse_thread && i == 0) {
      continue; // Skip over first/main thread when re-using the users thread, as this one will replace the first one.
    }
    memory.thread_for(i) = base::thread(&scheduler::work_thread_main_loop, &memory_.thread_state_for(i));
  }
}

scheduler::~scheduler() {
  terminate();
}

void scheduler::work_thread_main_loop() {
  auto &scheduler = thread_state::get().get_scheduler();
  while (true) {
    // Wait to be triggered
    scheduler.sync_barrier_.wait();

    // Check for shutdown
    if (scheduler.terminated_) {
      return;
    }

    scheduler.work_thread_work_section();

    // Sync back with main thread
    scheduler.sync_barrier_.wait();
  }
}

void scheduler::work_thread_work_section() {
  auto &my_state = thread_state::get();
  auto &my_task_manager = my_state.get_task_manager();

  auto const num_threads = my_state.get_scheduler().num_threads();
  auto const my_id = my_state.get_id();

  if (my_state.get_id() == 0) {
    // Main Thread, kick off by executing the user's main code block.
    main_thread_starter_function_->run();
  }

  while (!work_section_done_) {
    PLS_ASSERT(my_task_manager.check_task_chain(), "Must start stealing with a clean task chain.");

    // Steal Routine (will be continuously executed when there are no more fall through's).
    // TODO: move into separate function
    const size_t offset = my_state.get_rand() % num_threads;
    const size_t max_tries = num_threads;
    for (size_t i = 0; i < max_tries; i++) {
      // Perform steal
      size_t target = (offset + i) % num_threads;
      auto &target_state = my_state.get_scheduler().thread_state_for(target);
      bool steal_success = target_state.get_task_manager().steal_task(my_task_manager);

      if (steal_success) {
        // The stealing procedure correctly changed our chain and active task.
        // Now we need to perform the 'post steal' actions (manage resources and execute the stolen task).
        PLS_ASSERT(my_task_manager.check_task_chain_forward(&my_task_manager.get_active_task()),
                   "We are sole owner of this chain, it has to be valid!");

        // Move the traded in resource of this active task over to the stack of resources.
        auto *stolen_task = &my_task_manager.get_active_task();
        traded_cas_field stolen_task_cas = stolen_task->external_trading_deque_cas_.load();
        if (stolen_task_cas.is_filled_with_object()) {
          // Push the traded in resource on the resource stack to clear the traded_field for later steals/spawns.
          auto *exchanged_task = stolen_task_cas.get_trade_object();
          my_task_manager.push_resource_on_task(stolen_task, exchanged_task);

          traded_cas_field empty_field;
          traded_cas_field expected_field;
          expected_field.fill_with_trade_object(exchanged_task);
          if (stolen_task->external_trading_deque_cas_.compare_exchange_strong(expected_field, empty_field)) {
            // All good, nothing more to do
          } else {
            // The last other active thread took it as its spare resource...
            // ...remove our traded object from the stack again (it must be empty now and no one must access it anymore).
            PLS_ASSERT(expected_field.is_empty(),
                       "Must be empty, as otherwise no one will steal the 'spare traded task'.");

            auto current_root = stolen_task->resource_stack_root_.load();
            current_root.stamp++;
            current_root.value = 0;
            stolen_task->resource_stack_root_.store(current_root);
          }
        }

        // Execute the stolen task by jumping to it's continuation.
        PLS_ASSERT(stolen_task->continuation_.valid(),
                   "A task that we can steal must have a valid continuation for us to start working.");
        context_switcher::switch_context(std::move(stolen_task->continuation_));

        // ...now we are done with this steal attempt, loop over.
        break;
      }
    }
//    if (!my_cont_manager.falling_through()) {
//      base::this_thread::sleep(5);
//    }
  }
}

void scheduler::terminate() {
  if (terminated_) {
    return;
  }

  terminated_ = true;
  sync_barrier_.wait();

  for (unsigned int i = 0; i < num_threads_; i++) {
    if (reuse_thread_ && i == 0) {
      continue;
    }
    memory_.thread_for(i).join();
  }
}

thread_state &scheduler::thread_state_for(size_t id) { return memory_.thread_state_for(id); }

void scheduler::sync() {
  thread_state::get().get_task_manager().sync();
}

}
}
}