#include "pls/internal/helpers/profiler.h"

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

namespace pls {
namespace internal {
namespace scheduling {

task::task() :
    ref_count_{0},
    parent_{nullptr},
    deque_state_{0} {}

task::task(const task &other) :
    ref_count_{0},
    parent_{other.parent_},
    deque_state_{other.deque_state_} {}

void task::execute() {
  {
    PROFILE_WORK_BLOCK("execute task")
    auto last_executing = thread_state::get()->current_task_;
    thread_state::get()->current_task_ = this;

    execute_internal();

    thread_state::get()->current_task_ = last_executing;
  }

  wait_for_all();

  if (parent_ != nullptr) {
    parent_->ref_count_--;
  }
}

void task::wait_for_all() {
  while (ref_count_ > 0) {
    PROFILE_STEALING("get local sub task")
    task *local_task = thread_state::get()->scheduler_->get_local_task();
    PROFILE_END_BLOCK
    if (local_task != nullptr) {
      local_task->execute();
    } else {
      // Try to steal work.
      // External steal will be executed implicitly if success
      PROFILE_STEALING("steal work")
      task *stolen_task = thread_state::get()->scheduler_->steal_task();
      PROFILE_END_BLOCK
      if (stolen_task != nullptr) {
        stolen_task->execute();
      }
    }
  }
  tbb_task_->deque_.release_memory_until(deque_state_);
}

bool task::internal_stealing(abstract_task *other_task) {
  PROFILE_STEALING("task::internal_stealin")
  auto cast_other_task = reinterpret_cast<task *>(other_task);

  auto stolen_sub_task = cast_other_task->get_stolen_sub_task();
  if (stolen_sub_task == nullptr) {
    return false;
  } else {
    // Make sub-task belong to our task instance
    stolen_sub_task->tbb_task_ = this;
    stolen_sub_task->deque_state_ = deque_.save_state();
    // We will execute this next without explicitly moving it onto our stack storage
    last_stolen_ = stolen_sub_task;

    return true;
  }
}

bool task::split_task(base::swmr_spin_lock *lock) {
  PROFILE_STEALING("task::split_task")
  fork_join_sub_task *stolen_sub_task = get_stolen_sub_task();
  if (stolen_sub_task == nullptr) {
    return false;
  }
  task task{stolen_sub_task, this->unique_id()};

  // In success case, unlock.
  lock->reader_unlock();

  scheduler::execute_task(task, depth());
  return true;
}

void task::execute() {
  PROFILE_WORK_BLOCK("execute task");

  // Bind this instance to our OS thread
  // TODO: See if we did this right
  // my_stack_ = base::this_thread::state<thread_state>()->task_stack_;
  deque_.reset_base_pointer();

  root_task_->tbb_task_ = this;
  root_task_->deque_state_ = deque_.save_state();

  // Execute it on our OS thread until its finished
  root_task_->execute();
}

}
}
}