#include "callback_manager.h"

#include <memory>

#include <utility>

#include <vector>

#include "absl/functional/any_invocable.h"

#include "absl/synchronization/mutex.h"

namespace safepower_agent {

CallbackManager::Handle::Handle(std::shared_ptr<HandleImpl> impl)

    : impl_(std::move(impl)) {}

bool CallbackManager::Handle::pending() const {

  if (impl_ == nullptr) return false;

  absl::MutexLock lock(&impl_->mutex);

  return impl_->state == State::kPending || impl_->state == State::kExecuting;

}

bool CallbackManager::Handle::executed() const {

  if (impl_ == nullptr) return false;

  absl::MutexLock lock(&impl_->mutex);

  return impl_->state == State::kExecuted;

}

bool CallbackManager::Handle::TryCancel() {

  if (impl_ == nullptr) return false;

  absl::MutexLock lock(&impl_->mutex);

  if (impl_->state == State::kPending) {

    impl_->callback = nullptr;

    impl_->state = State::kCancelled;

    return true;

  }

  return false;

}

void CallbackManager::Handle::Wait() {

  if (impl_ == nullptr) return;

  absl::MutexLock lock(&impl_->mutex);

  impl_->mutex.Await(absl::Condition(

      +[](State* state) {

        return *state == State::kCancelled || *state == State::kExecuted;

      },

      &impl_->state));

}

CallbackManager::Handle CallbackManager::RunFirst(Callback callback) {

  auto impl = std::make_shared<HandleImpl>();

  impl->callback = std::move(callback);

  {

    absl::MutexLock lock(&mutex_);

    run_first_.push_back(impl);

  }

  return Handle(std::move(impl));

}

CallbackManager::Handle CallbackManager::RunLast(Callback callback) {

  auto impl = std::make_shared<HandleImpl>();

  impl->callback = std::move(callback);

  {

    absl::MutexLock lock(&mutex_);

    run_last_.push_back(impl);

  }

  return Handle(std::move(impl));

}

void CallbackManager::RunCallbackIfValid(

    const std::weak_ptr<HandleImpl>& weak_handle) {

  Callback cb;

  std::shared_ptr<HandleImpl> handle = weak_handle.lock();

  if (handle) {

    {

      absl::MutexLock lock(&handle->mutex);

      if (handle->state == State::kPending) {

        cb = std::exchange(handle->callback, nullptr);

        if (!cb) {

          // If the callback was registered as empty/nullptr, transition

          // directly to kExecuted so that we don't get stuck in kExecuting.

          handle->state = State::kExecuted;

          return;

        }

        handle->state = State::kExecuting;

      }

    }

  }

  if (cb) {

    std::move(cb)();

    absl::MutexLock lock(&handle->mutex);

    handle->state = State::kExecuted;

  }

}

void CallbackManager::RunCallbacks() {

  std::vector<std::weak_ptr<HandleImpl>> first_snapshot;

  std::vector<std::weak_ptr<HandleImpl>> last_snapshot;

  {

    absl::MutexLock lock(&mutex_);

    first_snapshot = std::move(run_first_);

    last_snapshot = std::move(run_last_);

    run_first_.clear();

    run_last_.clear();

  }

  // Run first callbacks in reverse (LIFO) order

  for (auto it = first_snapshot.rbegin(); it != first_snapshot.rend(); ++it) {

    RunCallbackIfValid(*it);

  }

  // Run last callbacks in forward (FIFO) order

  for (auto it = last_snapshot.begin(); it != last_snapshot.end(); ++it) {

    RunCallbackIfValid(*it);

  }

}

}  // namespace safepower_agent