From c817f3189f65f24da151bbc9bd0a8dba272176b2 Mon Sep 17 00:00:00 2001
From: gabest11 <gabest11@gmail.com>
Date: Mon, 4 Nov 2024 02:56:59 +0100
Subject: [PATCH] deep_sleep guard flag, to prevent sleep no matter what

---
 .../deep_sleep/deep_sleep_component.cpp       |  8 ++++-
 .../deep_sleep/deep_sleep_component.h         |  7 ++++
 .../components/deep_sleep/switch/__init__.py  | 36 +++++++++++++++++++
 esphome/components/deep_sleep/switch/switch.h | 18 ++++++++++
 4 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 esphome/components/deep_sleep/switch/__init__.py
 create mode 100644 esphome/components/deep_sleep/switch/switch.h

diff --git a/esphome/components/deep_sleep/deep_sleep_component.cpp b/esphome/components/deep_sleep/deep_sleep_component.cpp
index 1e7637f3e5..2eab0b831c 100644
--- a/esphome/components/deep_sleep/deep_sleep_component.cpp
+++ b/esphome/components/deep_sleep/deep_sleep_component.cpp
@@ -20,6 +20,10 @@ void DeepSleepComponent::setup() {
   } else {
     ESP_LOGD(TAG, "Not scheduling Deep Sleep, as no run duration is configured.");
   }
+
+  if (this->guard_switch_ != nullptr) {
+    this->guard_switch_->publish_state(this->guard_);
+  }
 }
 
 void DeepSleepComponent::dump_config() {
@@ -48,7 +52,7 @@ void DeepSleepComponent::set_sleep_duration(uint32_t time_ms) { this->sleep_dura
 void DeepSleepComponent::set_run_duration(uint32_t time_ms) { this->run_duration_ = time_ms; }
 
 void DeepSleepComponent::begin_sleep(bool manual) {
-  if (this->prevent_ && !manual) {
+  if (this->prevent_ && !manual || this->guard_) {
     this->next_enter_deep_sleep_ = true;
     return;
   }
@@ -72,5 +76,7 @@ void DeepSleepComponent::prevent_deep_sleep() { this->prevent_ = true; }
 
 void DeepSleepComponent::allow_deep_sleep() { this->prevent_ = false; }
 
+void DeepSleepComponent::set_guard(bool value) { this->guard_ = value; }
+
 }  // namespace deep_sleep
 }  // namespace esphome
diff --git a/esphome/components/deep_sleep/deep_sleep_component.h b/esphome/components/deep_sleep/deep_sleep_component.h
index 7a640b9ea5..316fb79fe9 100644
--- a/esphome/components/deep_sleep/deep_sleep_component.h
+++ b/esphome/components/deep_sleep/deep_sleep_component.h
@@ -14,6 +14,8 @@
 #include "esphome/core/time.h"
 #endif
 
+#include "esphome/components/switch/switch.h"
+
 #include <cinttypes>
 
 namespace esphome {
@@ -64,6 +66,8 @@ template<typename... Ts> class PreventDeepSleepAction;
  */
 class DeepSleepComponent : public Component {
  public:
+  SUB_SWITCH(guard)
+
   /// Set the duration in ms the component should sleep once it's in deep sleep mode.
   void set_sleep_duration(uint32_t time_ms);
 #if defined(USE_ESP32)
@@ -103,6 +107,8 @@ class DeepSleepComponent : public Component {
   void prevent_deep_sleep();
   void allow_deep_sleep();
 
+  void set_guard(bool value);
+
  protected:
   // Returns nullopt if no run duration is set. Otherwise, returns the run
   // duration before entering deep sleep.
@@ -127,6 +133,7 @@ class DeepSleepComponent : public Component {
   optional<uint32_t> run_duration_;
   bool next_enter_deep_sleep_{false};
   bool prevent_{false};
+  bool guard_{false};
 };
 
 extern bool global_has_deep_sleep;  // NOLINT(cppcoreguidelines-avoid-non-const-global-variables)
diff --git a/esphome/components/deep_sleep/switch/__init__.py b/esphome/components/deep_sleep/switch/__init__.py
new file mode 100644
index 0000000000..d7bf8ee496
--- /dev/null
+++ b/esphome/components/deep_sleep/switch/__init__.py
@@ -0,0 +1,36 @@
+import esphome.codegen as cg
+from esphome.components import switch
+import esphome.config_validation as cv
+from esphome.const import (
+    DEVICE_CLASS_SWITCH,
+)
+from .. import (
+    DeepSleepComponent,
+    deep_sleep_ns,
+)
+
+CONF_DEEP_SLEEP_ID = "deep_sleep_id"
+CONF_GUARD = "guard"
+ICON_SLEEP = "mdi:sleep"
+
+GuardSwitch = deep_sleep_ns.class_("GuardSwitch", switch.Switch)
+
+CONFIG_SCHEMA = cv.Schema(
+    {
+        cv.GenerateID(CONF_DEEP_SLEEP_ID): cv.use_id(DeepSleepComponent),
+        cv.Optional(CONF_GUARD): switch.switch_schema(
+            GuardSwitch,
+            device_class=DEVICE_CLASS_SWITCH,
+            icon=ICON_SLEEP,
+        ),
+    }
+)
+
+
+async def to_code(config):
+    parent = await cg.get_variable(config[CONF_DEEP_SLEEP_ID])
+
+    if CONF_GUARD in config:
+        s = await switch.new_switch(config[CONF_GUARD])
+        await cg.register_parented(s, parent)
+        cg.add(parent.set_guard_switch(s))
diff --git a/esphome/components/deep_sleep/switch/switch.h b/esphome/components/deep_sleep/switch/switch.h
new file mode 100644
index 0000000000..4a1cdaa253
--- /dev/null
+++ b/esphome/components/deep_sleep/switch/switch.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include "esphome/components/switch/switch.h"
+#include "../deep_sleep_component.h"
+
+namespace esphome {
+namespace deep_sleep {
+
+class GuardSwitch : public switch_::Switch, public Parented<DeepSleepComponent> {
+ protected:
+  void write_state(bool value) override {
+    this->publish_state(value);
+    this->parent_->set_guard(value);
+  }
+};
+
+}  // namespace deep_sleep
+}  // namespace esphome
\ No newline at end of file