From 176304db46cff83dbec40d1b36d56a7f717dfaf8 Mon Sep 17 00:00:00 2001
From: JonatanRek <vasek@steelants.cz>
Date: Mon, 3 Mar 2025 09:46:15 +0100
Subject: [PATCH] Fixes for actions

---
 __init__.py    | 22 ++++++++++++++++++----
 automation.cpp | 18 ------------------
 automation.h   | 34 +++++++++++++++++++++++++++++++---
 3 files changed, 49 insertions(+), 25 deletions(-)
 delete mode 100644 automation.cpp

diff --git a/__init__.py b/__init__.py
index 1664827..a944259 100644
--- a/__init__.py
+++ b/__init__.py
@@ -22,6 +22,7 @@ DEPENDENCIES = ['uart']
 CONF_MARLIN2_ID = "marlin2_id"
 
 Marlin2 = cg.esphome_ns.class_('Marlin2', cg.Component)
+Marlin2WriteAction = uart_ns.class_("Marlin2WriteAction", automation.Action)
 
 CONFIG_SCHEMA = cv.All(
     cv.Schema({
@@ -31,10 +32,23 @@ CONFIG_SCHEMA = cv.All(
     .extend(uart.UART_DEVICE_SCHEMA),
 )
 
-@automation.register_action("marlin2.set_bed_temperature", SetTemperatureAction, SWITCH_ACTION_SCHEMA)
-async def marlin2_set_bed_temperature_to_code(config, action_id, template_arg, args):
-    paren = await cg.get_variable(config[CONF_ID])
-    return cg.new_Pvariable(action_id, template_arg, paren)
+MARLIN2_ACTION_SCHEMA = maybe_simple_id(
+    {
+        cv.Required(CONF_ID): cv.use_id(Marlin2),
+    }
+)
+
+@automation.register_action(
+    "marlin2.write",
+    Marlin2WriteAction,
+    cv.maybe_simple_value(
+        {
+            cv.GenerateID(): cv.use_id(UARTComponent),
+            cv.Required(CONF_DATA): cv.templatable(validate_raw_data),
+        },
+        key=CONF_DATA,
+    ),
+)
 
 async def to_code(config):
     var = cg.new_Pvariable(config[CONF_ID])
diff --git a/automation.cpp b/automation.cpp
deleted file mode 100644
index a2afc3d..0000000
--- a/automation.cpp
+++ /dev/null
@@ -1,18 +0,0 @@
-#pragma once
-
-#include "esphome/core/component.h"
-#include "esphome/core/automation.h"
-#include "esphome/components/switch/switch.h"
-
-namespace esphome {
-
-template<typename... Ts> class SetTemperatureAction : public Action<Ts...> {
-    public:
-        explicit SetTemperatureAction(Marlin2 *a_marlin2) : switch_(a_marlin2) {}
-
-    void play(Ts... x) override { this->switch_->turn_on(); }
-
-    protected:
-        Switch *switch_;
-};
-}  // namespace esphome
\ No newline at end of file
diff --git a/automation.h b/automation.h
index f95c0b0..ee4716d 100644
--- a/automation.h
+++ b/automation.h
@@ -1,8 +1,36 @@
-#include "automation.h"
-#include "esphome/core/log.h"
+#pragma once
+
+#include "uart.h"
+#include "esphome/core/automation.h"
+
+#include <vector>
 
 namespace esphome {
 
-static const char *const TAG = "switch.automation";
+template<typename... Ts> class Marlin2WriteAction : public Action<Ts...>, public Parented<UARTComponent> {
+    public:
+        void set_data_template(std::function<std::vector<uint8_t>(Ts...)> func) {
+            this->data_func_ = func;
+            this->static_ = false;
+        }
+        void set_data_static(const std::vector<uint8_t> &data) {  
+            this->data_static_ = data;
+            this->static_ = true;
+        }
+
+        void play(Ts... x) override {
+            if (this->static_) {
+                this->parent_->write_array(this->data_static_);
+            } else {
+                auto val = this->data_func_(x...);
+                this->parent_->write_array(val);
+            }
+        }
+
+    protected:
+        bool static_{false};
+        std::function<std::vector<uint8_t>(Ts...)> data_func_{};
+        std::vector<uint8_t> data_static_{};
+};
 
 }  // namespace esphome
\ No newline at end of file