#include <catch.hpp>

#include <pls/pls.h>

using namespace pls;

class once_sub_task: public tbb_sub_task {
    std::atomic<int>* counter_;
    int children_;

protected:
    void execute_internal() override {
        (*counter_)++;
        for (int i = 0; i < children_; i++) {
            spawn_child(once_sub_task(counter_, children_ - 1));
        }
    }

public:
    explicit once_sub_task(std::atomic<int>* counter, int children):
        tbb_sub_task(),
        counter_{counter},
        children_{children} {}
};

class force_steal_sub_task: public tbb_sub_task {
    std::atomic<int>* parent_counter_;
    std::atomic<int>* overall_counter_;

protected:
    void execute_internal() override {
        (*overall_counter_)--;
        if (overall_counter_->load() > 0) {
            std::atomic<int> counter{1};
            spawn_child(force_steal_sub_task(&counter, overall_counter_));
            while (counter.load() > 0)
                ; // Spin...
        }

        (*parent_counter_)--;
    }

public:
    explicit force_steal_sub_task(std::atomic<int>* parent_counter, std::atomic<int>* overall_counter):
            tbb_sub_task(),
            parent_counter_{parent_counter},
            overall_counter_{overall_counter} {}
};

TEST_CASE( "tbb task are scheduled correctly", "[internal/scheduling/tbb_task.h]") {
    static static_scheduler_memory<8, 2 << 12> my_scheduler_memory;

    SECTION("tasks are executed exactly once") {
        scheduler my_scheduler{&my_scheduler_memory, 2};
        int start_counter = 4;
        int total_tasks = 1 + 4 + 4 * 3 + 4 * 3 * 2 + 4 * 3 * 2 * 1;
        std::atomic<int> counter{0};

        my_scheduler.perform_work([&] (){
            once_sub_task sub_task{&counter, start_counter};
            tbb_task task{&sub_task};
            scheduler::execute_task(task);
        });

        REQUIRE(counter.load() == total_tasks);
        my_scheduler.terminate(true);
    }

    SECTION("tasks can be stolen") {
        scheduler my_scheduler{&my_scheduler_memory, 8};
        my_scheduler.perform_work([&] (){
            std::atomic<int> dummy_parent{1}, overall_counter{8};
            force_steal_sub_task sub_task{&dummy_parent, &overall_counter};
            tbb_task task{&sub_task};
            scheduler::execute_task(task);
        });
        my_scheduler.terminate(true);
    }
}