#pragma once

#include "esphome/core/component.h"
#include "esphome/core/defines.h"
#include "esphome/core/helpers.h"
#include "esphome/core/preferences.h"

#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
#include <esp_ota_ops.h>
#endif

namespace esphome::safe_mode {

/// RTC key for storing boot loop counter - used by safe_mode and preferences backends
constexpr uint32_t RTC_KEY = 233825507UL;

/// SafeModeComponent provides a safe way to recover from repeated boot failures
class SafeModeComponent final : public Component {
 public:
  bool should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_time, uint32_t boot_is_good_after);

  /// Set to true if the next startup will enter safe mode
  void set_safe_mode_pending(const bool &pending);
  bool get_safe_mode_pending();

  void dump_config() override;
  float get_setup_priority() const override;
  void loop() override;

  void clean_rtc();

  void on_safe_shutdown() override;

  void mark_successful();

#ifdef USE_SAFE_MODE_CALLBACK
  template<typename F> void add_on_safe_mode_callback(F &&callback) {
    this->safe_mode_callback_.add(std::forward<F>(callback));
  }
#endif

 protected:
  void write_rtc_(uint32_t val);
  uint32_t read_rtc_();

  // Group all 4-byte aligned members together to avoid padding
  uint32_t safe_mode_boot_is_good_after_{60000};  ///< The amount of time after which the boot is considered successful
  uint32_t safe_mode_enable_time_{60000};         ///< The time safe mode should remain active for
  uint32_t safe_mode_rtc_value_{0};
  uint32_t safe_mode_start_time_{0};  ///< stores when safe mode was enabled
#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
  esp_ota_img_states_t ota_state_{ESP_OTA_IMG_UNDEFINED};  // 4-byte enum
#endif
  // Group 1-byte members together to minimize padding
  bool boot_successful_{false};  ///< set to true after boot is considered successful
  uint8_t safe_mode_num_attempts_{0};
#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
  bool app_ota_possible_{true};
#endif
  // Larger objects at the end
  ESPPreferenceObject rtc_;
#ifdef USE_SAFE_MODE_CALLBACK
  StaticCallbackManager<ESPHOME_SAFE_MODE_CALLBACK_COUNT, void()> safe_mode_callback_{};
#endif

  static const uint32_t ENTER_SAFE_MODE_MAGIC =
      0x5afe5afe;  ///< a magic number to indicate that safe mode should be entered on next boot
};

}  // namespace esphome::safe_mode
