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

#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() {
                PROFILE_STEALING("abstract_task::steal_work")
                const auto my_state = base::this_thread::state<thread_state>();
                const auto my_scheduler = my_state->scheduler_;

                const size_t my_id = my_state->id_;
                const size_t offset = my_state->random_() % my_scheduler->num_threads();
                const size_t max_tries = 1; // my_scheduler->num_threads(); TODO: Tune this value
                for (size_t i = 0; i < max_tries; i++) {
                    size_t target = (offset + i) % my_scheduler->num_threads();
                    if (target == my_id) {
                        continue;
                    }
                    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
                    PROFILE_STEALING("Go 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_;
                    }
                    PROFILE_END_BLOCK

                    // Try to steal 'internal', e.g. for_join_sub_tasks in a fork_join_task constellation
                    PROFILE_STEALING("Internal Steal")
                    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_;
                        }
                    }
                    PROFILE_END_BLOCK;


                    // Execute 'top level task steal' if possible
                    // (only try deeper tasks to keep depth restricted stealing).
                    PROFILE_STEALING("Top Level Steal")
                    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_;
                    }
                    PROFILE_END_BLOCK;
                    target_state->lock_.unlock();
                }

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