
#ifndef PLS_RUN_ON_N_THREADS_TASK_H
#define PLS_RUN_ON_N_THREADS_TASK_H

#include <mutex>

#include "pls/internal/base/spin_lock.h"
#include "pls/internal/base/thread.h"

#include "abstract_task.h"
#include "thread_state.h"
#include "scheduler.h"

namespace pls {
namespace internal {
namespace scheduling {

template<typename Function>
class run_on_n_threads_task : public abstract_task {
  template<typename F>
  friend
  class run_on_n_threads_task_worker;

  Function function_;

  // Improvement: Remove lock and replace by atomic variable (performance)
  int counter;
  base::spin_lock counter_lock_;

  int decrement_counter() {
    std::lock_guard<base::spin_lock> lock{counter_lock_};
    counter--;
    return counter;
  }

  int get_counter() {
    std::lock_guard<base::spin_lock> lock{counter_lock_};
    return counter;
  }
 public:
  static constexpr auto create_id = helpers::unique_id::create<run_on_n_threads_task<Function>>;

  run_on_n_threads_task(Function function, int num_threads) :
      abstract_task{0, create_id()},
      function_{function},
      counter{num_threads - 1} {}

  void execute() override {
    // Execute our function ONCE
    function_();

    // Steal until we are finished (other threads executed)
    do {
      steal_work();
    } while (get_counter() > 0);

    std::cout << "Finished Master!" << std::endl;
  }

  bool internal_stealing(abstract_task * /*other_task*/) override {
    return false;
  }

  bool split_task(base::swmr_spin_lock *lock) override;
};

template<typename Function>
class run_on_n_threads_task_worker : public abstract_task {
  Function function_;
  run_on_n_threads_task<Function> *root_;
 public:
  static constexpr auto create_id = helpers::unique_id::create<run_on_n_threads_task_worker<Function>>;

  run_on_n_threads_task_worker(Function function, run_on_n_threads_task<Function> *root) :
      abstract_task{0, create_id()},
      function_{function},
      root_{root} {}

  void execute() override {
    if (root_->decrement_counter() >= 0) {
      function_();
      std::cout << "Finished Worker!" << std::endl;
    } else {
      std::cout << "Abandoned Worker!" << std::endl;
    }
  }

  bool internal_stealing(abstract_task * /*other_task*/) override {
    return false;
  }

  bool split_task(base::swmr_spin_lock * /*lock*/) override {
    return false;
  }
};

template<typename Function>
bool run_on_n_threads_task<Function>::split_task(base::swmr_spin_lock *lock) {
  if (get_counter() <= 0) {
    return false;
  }
  // In success case, unlock.
  lock->reader_unlock();

  auto scheduler = base::this_thread::state<thread_state>()->scheduler_;
  auto task = run_on_n_threads_task_worker<Function>{function_, this};
  scheduler->execute_task(task, depth());
  return true;
}

template<typename Function>
run_on_n_threads_task<Function> create_run_on_n_threads_task(Function function, int num_threads) {
  return run_on_n_threads_task<Function>{function, num_threads};
}

}
}
}

#endif //PLS_RUN_ON_N_THREADS_TASK_H
