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

namespace pls {
    namespace internal {
        namespace scheduling {
            bool abstract_task::steal_work() {
                auto my_state = base::this_thread::state<thread_state>();
                auto my_scheduler = my_state->scheduler_;

                int my_id = my_state->id_;
                for (size_t i = 1; i < my_scheduler->num_threads(); i++) {
                    size_t target = (my_id + i) % my_scheduler->num_threads();
                    auto target_state = my_scheduler->thread_state_for(target);

                    // TODO: Cleaner Locking Using std::guarded_lock
                    target_state->lock_.lock();

                    // Dig down to our level
                    abstract_task* current_task = target_state->root_task_;
                    while (current_task != nullptr && current_task->depth() < depth()) {
                        current_task = current_task->child_task_;
                    }

                    if (current_task != nullptr) {
                        // See if it equals our type and depth of task
                        if (current_task->unique_id_ == unique_id_ &&
                            current_task->depth_ == depth_) {
                            if (internal_stealing(current_task)) {
                                // internal steal was a success, hand it back to the internal scheduler
                                target_state->lock_.unlock();
                                return true;
                            }

                            // No success, we need to steal work from a deeper level using 'top level task stealing'
                            current_task = current_task->child_task_;
                        }
                    }


                    // Execute 'top level task steal' if possible
                    // (only try deeper tasks to keep depth restricted stealing)
                    while (current_task != nullptr) {
                        auto lock = &target_state->lock_;
                        if (current_task->split_task(lock)) {
                            // internal steal was no success (we did a top level task steal)
                            return false;
                        }

                        current_task = current_task->child_task_;
                    }
                    target_state->lock_.unlock();
                }

                // internal steal was no success
                return false;
            };
        }
    }
}