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

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

namespace pls {
namespace internal {
namespace scheduling {

fork_join_sub_task::fork_join_sub_task() :
    ref_count_{0},
    parent_{nullptr},
    tbb_task_{nullptr},
    deque_state_{0} {}

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

void fork_join_sub_task::execute() {
  PROFILE_WORK_BLOCK("execute sub_task")
  auto last_executing = tbb_task_->currently_executing_;
  tbb_task_->currently_executing_ = this;
  execute_internal();
  tbb_task_->currently_executing_ = last_executing;
  PROFILE_END_BLOCK
  wait_for_all();

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

void fork_join_sub_task::wait_for_all() {
  while (ref_count_ > 0) {
    PROFILE_STEALING("get local sub task")
    fork_join_sub_task *local_task = tbb_task_->get_local_sub_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")
      bool internal_steal_success = tbb_task_->steal_work();
      PROFILE_END_BLOCK
      if (internal_steal_success) {
        tbb_task_->last_stolen_->execute();
      }
    }
  }
  tbb_task_->deque_.release_memory_until(deque_state_);
}

fork_join_sub_task *fork_join_task::get_local_sub_task() {
  return deque_.pop_tail();
}

fork_join_sub_task *fork_join_task::get_stolen_sub_task() {
  return deque_.pop_head();
}

fork_join_sub_task *fork_join_sub_task::current() {
  return dynamic_cast<fork_join_task *>(scheduler::current_task())->currently_executing();
}

bool fork_join_task::internal_stealing(abstract_task *other_task) {
  PROFILE_STEALING("fork_join_task::internal_stealin")
  auto cast_other_task = reinterpret_cast<fork_join_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 fork_join_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 fork_join_task::split_task(base::swmr_spin_lock *lock) {
  PROFILE_STEALING("fork_join_task::split_task")
  fork_join_sub_task *stolen_sub_task = get_stolen_sub_task();
  if (stolen_sub_task == nullptr) {
    return false;
  }
  fork_join_task task{stolen_sub_task, this->unique_id()};

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

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

void fork_join_task::execute() {
  PROFILE_WORK_BLOCK("execute fork_join_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();
}

fork_join_sub_task *fork_join_task::currently_executing() const { return currently_executing_; }

fork_join_task::fork_join_task(fork_join_sub_task *root_task,
                               const abstract_task::id &id) :
    abstract_task{0, id},
    root_task_{root_task},
    currently_executing_{nullptr},
    deque_{base::this_thread::state<thread_state>()->task_stack_},
    last_stolen_{nullptr} {}

}
}
}