#ifndef PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_DISRUPTION_MANAGER_H_
#define PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_DISRUPTION_MANAGER_H_

#include <list>

#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "source_location"

namespace safepower_agent {
class DisruptionManager {
  struct PrivateToken {};

 public:
  using DisruptionCallback = absl::AnyInvocable<void() &&>;

  DisruptionManager() = default;
  DisruptionManager(const DisruptionManager&) = delete;
  DisruptionManager& operator=(const DisruptionManager&) = delete;
  DisruptionManager(DisruptionManager&&) = delete;
  DisruptionManager& operator=(DisruptionManager&&) = delete;

  class CallbackHandle {
   public:
    CallbackHandle() = default;
    CallbackHandle(PrivateToken, absl::Mutex& mutex,
                   std::list<DisruptionCallback>& list,
                   std::list<DisruptionCallback>::iterator it);
    CallbackHandle(CallbackHandle&& other);
    CallbackHandle& operator=(CallbackHandle&& other);
    ~CallbackHandle();

   private:
    absl::Mutex*  mutex_ = nullptr;
    std::list<DisruptionCallback>*  list_ = nullptr;
    std::list<DisruptionCallback>::iterator it_;
  };

  // Indicates that a disruption is expected to happen in the near future.
  // Returns an error if there is already a pending disruption.
  absl::Status ExpectDisruptionIn(absl::Duration timeout)
      ABSL_LOCKS_EXCLUDED(disruption_mutex_);
  // Cancels the pending disruption. Returns an error if there is no pending
  // disruption.
  absl::Status CancelDisruption() ABSL_LOCKS_EXCLUDED(disruption_mutex_);

  // Returns the remaining duration of the pending disruption, or zero if there
  // is no pending disruption.
  absl::Duration PendingDisruption() const
      ABSL_LOCKS_EXCLUDED(disruption_mutex_);

  // Runs a callback when a disruption is expected. The callback will be called
  // at most once. The move-only returned handle can be destroyed to cancel the
  // callback.
  // The callback may be used to delay the disruption by
  // ensuring that any required side effects have taken place before it returns.
  [[nodiscard]] CallbackHandle OnDisruptionStart(DisruptionCallback callback)
      ABSL_LOCKS_EXCLUDED(disruption_mutex_);
  // Runs a callback when a disruption is no longer expected because it was
  // cancelled or because there was a timeout. The callback will be called at
  // most once. The move-only returned handle can be destroyed to cancel the
  // callback.
  [[nodiscard]] CallbackHandle OnDisruptionEnd(DisruptionCallback callback)
      ABSL_LOCKS_EXCLUDED(disruption_mutex_);

 private:
  absl::Status RunCallbacks(
      std::list<DisruptionCallback>& list, absl::Duration delay,
      std::source_location location = std::source_location::current());
  CallbackHandle OnDisruptionStartHelper(DisruptionCallback callback)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(disruption_mutex_);

  CallbackHandle OnDisruptionEndHelper(DisruptionCallback callback)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(disruption_mutex_);

  mutable absl::Mutex disruption_mutex_;
  absl::Time expect_disruption_until_ ABSL_GUARDED_BY(disruption_mutex_) =
      absl::InfinitePast();
  std::list<DisruptionCallback> disruption_start_callbacks_
      ABSL_GUARDED_BY(disruption_mutex_);
  std::list<DisruptionCallback> disruption_end_callbacks_
      ABSL_GUARDED_BY(disruption_mutex_);
};
}  // namespace safepower_agent

#endif  // PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_DISRUPTION_MANAGER_H_