#ifndef PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_

#define PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_

#include <algorithm>

#include <array>

#include <cstddef>

#include <list>

#include <memory>

#include <string>

#include <utility>

#include <vector>

#include "callback_manager.h"

#include "state_merge.h"

#include "absl/base/nullability.h"

#include "absl/base/thread_annotations.h"

#include "absl/functional/any_invocable.h"

#include "absl/functional/function_ref.h"

#include "absl/log/log.h"

#include "absl/synchronization/mutex.h"

namespace safepower_agent {

template <typename StateProto>

class StateUpdater {

 public:

  class Listener;

  using ListenerIterator = typename std::list<Listener* >::iterator;

  class Listener {

    struct StateUpdaterRef {

      std::weak_ptr<StateUpdater> ptr;

      StateUpdater* raw_ptr;  // For compares in destructor.

      ListenerIterator itr;

    };

   public:

    explicit Listener(bool is_passive_listener = false)

        : is_passive_listener_(is_passive_listener) {}

    virtual ~Listener() {

      if (!state_updater_.empty()) {

        LOG(DFATAL) << "Listener destroyed without calling StopListening";

      }

      StopListening();

    }

    // Derived classes must explicitly call this at the start of their

    // destructor if they access any member variables or resources during

    // callbacks, since those resources will be destroyed before the

    // base ~Listener() runs.

    void StopListening() {

      for (const StateUpdaterRef& ref : state_updater_) {

        std::shared_ptr<StateUpdater> state_updater = ref.ptr.lock();

        if (state_updater != nullptr) {

          state_updater->RemoveListener(ref.itr);

        }

      }

      state_updater_.clear();

    }

    template <typename... Args>

    void Listen(const Args&... args) {

      static_assert(sizeof...(args) > 0);

      if (!state_updater_.empty()) {

        LOG(DFATAL) << "Listener::Listen called multiple times";

        return;

      }

      state_updater_.reserve(sizeof...(args));

      std::array done = {Insert(args)...};

      for (size_t i = 0; i < done.size(); ++i) {

        if (done[i]) {

          MaybeDone(*state_updater_[i].raw_ptr);

        }

      }

    }

    bool IsPassiveListener() const { return is_passive_listener_; }

    virtual void UpdateState([[maybe_unused]] const StateProto& previous_state,

                             [[maybe_unused]] const StateProto& update) {

      LOG(DFATAL) << "UpdateState called on base Listener";

    }

    void MaybeDone(const StateUpdater& updater) {

      for (StateUpdaterRef& ref : state_updater_) {

        if (ref.raw_ptr == &updater) {

          ref.raw_ptr = nullptr;

          ++done_count_;

          if (done_count_ >= state_updater_.size()) {

            Done();

          }

          return;

        }

      }

      LOG(DFATAL) << "Listener not listening to this updater";

    }

   protected:

    virtual void Done() { LOG(DFATAL) << "Done called on base Listener"; }

   private:

    bool Insert(const  std::shared_ptr<StateUpdater>& updater) {

      for (const StateUpdaterRef& ref : state_updater_) {

        if (ref.ptr.lock() == updater) {

          LOG(DFATAL) << "Listener already listening to this updater";

          return true;

        }

      }

      bool final = updater->AddListener(

          *this, [&](ListenerIterator itr) {

            state_updater_.push_back({updater, updater.get(), itr});

          });

      return final;

    }

    std::vector<StateUpdaterRef> state_updater_;

    size_t done_count_ = 0;

    bool is_passive_listener_ = false;

  };

  friend class Listener;

  explicit StateUpdater(StateProto initial_state = {}, bool final_state = false)

      : state_(std::move(initial_state)), final_(final_state) {}

  ~StateUpdater() {

    if (!listeners_mutex_.WriterTryLock()) {

      LOG(DFATAL) << "StateUpdater destroyed while mutex was held";

      return;

    }

    for (Listener* listener : listeners_) {

      listener->MaybeDone(*this);

    }

    listeners_.clear();

    listeners_mutex_.WriterUnlock();

  }

  // Returns a copy of the state. If you want to access the state without

  // copying (e.g. for performance reasons with large states), prefer using

  // WithState.

  StateProto state() const {

    absl::MutexLock lock(&state_mutex_);

    return state_;

  }

  template <typename Func>

  auto WithState(Func&& f) const ABSL_LOCKS_EXCLUDED(state_mutex_) {

    absl::MutexLock lock(&state_mutex_);

    return std::forward<Func>(f)(state_);

  }

  std::string DebugString() const ABSL_LOCKS_EXCLUDED(state_mutex_) {

    return WithState(

        [](const StateProto& state) { return state.DebugString(); });

  }

  void UpdateState(const StateProto& update, bool final_state = false)

      ABSL_LOCKS_EXCLUDED(state_mutex_, listeners_mutex_) {

    absl::MutexLock lock_state(&state_mutex_);

    final_ = final_state;

    absl::MutexLock lock_listeners(&listeners_mutex_);

    for (auto listener : listeners_) {

      listener->UpdateState(state_, update);

      if (final_state) {

        listener->MaybeDone(*this);

      }

    }

    ApplyStateUpdate(state_, update);

  }

  // Runs a callback when the updater has no more listeners. The callback will

  // be called immediately if the updater is already idle. The move-only

  // returned handle can be destroyed to cancel the callback.

  // Idle callbacks are run in LIFO order, while active callbacks are run in

  // FIFO order. This reduces the likelihood of deadlocks: For example, if a

  // resource is acquired when active and released when idle, another pair of

  // callbacks that are added later will always see that resource acquired.

  CallbackManager::Handle OnIdle(absl::AnyInvocable<void() &&> callback)

      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {

    if (IsIdle()) {

      std::move(callback)();

      return {};

    }

    return activity_callbacks_.RunFirst(std::move(callback));

  }

  // Runs a callback when the updater has at least one listener. The callback

  // will be called immediately if the updater already has listeners. See OnIdle

  // for more details.

  CallbackManager::Handle OnActive(absl::AnyInvocable<void() &&> callback)

      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {

    if (!IsIdle()) {

      std::move(callback)();

      return {};

    }

    return activity_callbacks_.RunLast(std::move(callback));

  }

 private:

  bool IsIdleInternal() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(listeners_mutex_) {

    for (const auto* listener : listeners_) {

      if (!listener->IsPassiveListener()) {

        return false;

      }

    }

    return true;

  }

  bool IsIdle() const ABSL_LOCKS_EXCLUDED(listeners_mutex_) {

    absl::MutexLock lock(&listeners_mutex_);

    return IsIdleInternal();

  }

  bool AddListener(Listener& listener,

                   absl::FunctionRef<void(ListenerIterator)> register_func)

      ABSL_LOCKS_EXCLUDED(state_mutex_, listeners_mutex_) {

    bool run_callbacks;

    bool result;

    {

      absl::MutexLock lock_state(&state_mutex_);

      listener.UpdateState({}, state_);

      absl::MutexLock lock_listeners(&listeners_mutex_);

      // IsIdleInternal() must be checked before push_front modifies the list

      // to correctly evaluate the pre-transition idle state.

      run_callbacks = IsIdleInternal() && !listener.IsPassiveListener();

      // push_front must happen before register_func so the iterator is valid.

      listeners_.push_front(&listener);

      register_func(listeners_.begin());

      result = final_;

    }

    if (run_callbacks) {

      activity_callbacks_.RunCallbacks();

    }

    return result;

  }

  void RemoveListener(ListenerIterator handle)

      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {

    bool becomes_idle = false;

    {

      absl::MutexLock lock(&listeners_mutex_);

      if (handle == listeners_.end()) return;

      bool was_idle = IsIdleInternal();

      listeners_.erase(handle);

      becomes_idle = !was_idle && IsIdleInternal();

    }

    if (becomes_idle) {

      activity_callbacks_.RunCallbacks();

    }

  }

  mutable absl::Mutex state_mutex_ ABSL_ACQUIRED_BEFORE(listeners_mutex_);

  mutable absl::Mutex listeners_mutex_;

  // This is a std::list so that we have stable iterators

  std::list<Listener* > listeners_

      ABSL_GUARDED_BY(listeners_mutex_);

  StateProto state_ ABSL_GUARDED_BY(state_mutex_);

  bool final_ ABSL_GUARDED_BY(state_mutex_);

  // Callbacks that are run when the updater becomes idle (no listeners) or

  // active (at least one listener).

  CallbackManager activity_callbacks_;

};

}  // namespace safepower_agent

#endif  // PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_