#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 "callback_manager.h"
#include "absl/base/thread_annotations.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "source_location"

namespace safepower_agent {
class DisruptionManager {
 public:
  using DisruptionCallback = absl::AnyInvocable<void() &&>;
  using CallbackHandle = CallbackManager::Handle;

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

  // 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);
  // 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);

 private:
  absl::Status RunCallbacks(
      CallbackManager& callback_manager, absl::string_view name,
      absl::Duration delay,
      std::source_location location = std::source_location::current());

  mutable absl::Mutex disruption_mutex_;
  absl::Time expect_disruption_until_ ABSL_GUARDED_BY(disruption_mutex_) =
      absl::InfinitePast();
  CallbackManager disruption_start_callbacks_;
  CallbackManager disruption_end_callbacks_;
};
}  // namespace safepower_agent

#endif  // PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_DISRUPTION_MANAGER_H_