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

namespace pls {
    namespace internal {
        namespace scheduling {
            tbb_sub_task::tbb_sub_task():
                base::deque_item{},
                ref_count_{0},
                parent_{nullptr},
                tbb_task_{nullptr},
                stack_state_{nullptr} {}

            tbb_sub_task::tbb_sub_task(const tbb_sub_task& other): base::deque_item(other) {
                // Do Nothing, will be inited after this anyways
            }

            void tbb_sub_task::execute()  {
                execute_internal();
                wait_for_all();

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

            void tbb_sub_task::spawn_child_internal(tbb_sub_task* sub_task) {
                // Keep our refcount up to date
                ref_count_++;

                // Assign forced values
                sub_task->parent_ = this;
                sub_task->tbb_task_ = tbb_task_;
                sub_task->stack_state_ = tbb_task_->my_stack_->save_state();

                tbb_task_->deque_.push_tail(sub_task);
            }

            void tbb_sub_task::wait_for_all() {
                while (ref_count_ > 0) {
                    tbb_sub_task* local_task = tbb_task_->get_local_sub_task();
                    if (local_task != nullptr) {
                        local_task->execute();
                    } else {
                        // Try to steal work.
                        // External steal will be executed implicitly if success
                        if (tbb_task_->steal_work()) {
                            tbb_task_->last_stolen_->execute();
                        }
                    }
                }
                tbb_task_->my_stack_->reset_state(stack_state_);
            }

            tbb_sub_task* tbb_task::get_local_sub_task() {
                return deque_.pop_tail();
            }

            tbb_sub_task* tbb_task::get_stolen_sub_task() {
                return deque_.pop_head();
            }

            bool tbb_task::internal_stealing(abstract_task* other_task) {
                auto cast_other_task = reinterpret_cast<tbb_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 tbb_task instance
                    stolen_sub_task->tbb_task_ = this;
                    stolen_sub_task->stack_state_ = my_stack_->save_state();
                    // We will execute this next without explicitly moving it onto our stack storage
                    last_stolen_ = stolen_sub_task;

                    return true;
                }
            }

            bool tbb_task::split_task(base::spin_lock* lock) {
                tbb_sub_task* stolen_sub_task = get_stolen_sub_task();
                if (stolen_sub_task == nullptr) {
                    return false;
                }
                tbb_task task{stolen_sub_task};

                // In success case, unlock.
                // TODO: this locking is complicated and error prone.
                lock->unlock();

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