From 313efb7604e012d962aa192163dd103967d4d51e Mon Sep 17 00:00:00 2001
From: Djuri Baars <dsbaars@users.noreply.github.com>
Date: Sat, 8 Jun 2024 15:07:15 +0200
Subject: [PATCH] Better handling of steal screen, add firmware web update
 functionality

---
 src/lib/block_notify.cpp |  1 +
 src/lib/webserver.cpp    | 86 ++++++++++++++++++++++++++++++++++++++++
 src/lib/webserver.hpp    |  4 ++
 3 files changed, 91 insertions(+)

diff --git a/src/lib/block_notify.cpp b/src/lib/block_notify.cpp
index ab546b0..83a3be3 100644
--- a/src/lib/block_notify.cpp
+++ b/src/lib/block_notify.cpp
@@ -192,6 +192,7 @@ void onWebsocketBlockMessage(esp_websocket_event_data_t *event_data)
           esp_timer_start_periodic(screenRotateTimer,
                                    timerPeriod * usPerSecond);
         }
+        vTaskDelay(pdMS_TO_TICKS(315*NUM_SCREENS)); // Extra delay because of screen switching
       }
 
       if (preferences.getBool("ledFlashOnUpd", false))
diff --git a/src/lib/webserver.cpp b/src/lib/webserver.cpp
index 7b73823..d242eb7 100644
--- a/src/lib/webserver.cpp
+++ b/src/lib/webserver.cpp
@@ -67,6 +67,9 @@ void setupWebserver()
   // server.on("^\\/api\\/lights\\/([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$", HTTP_GET,
   // onApiLightsSetColor);
 
+  server.on("/firmware/update", HTTP_POST, onFirmwareUpdate, asyncFirmwareUpdateHandler);
+  server.on("/firmware/update_webui", HTTP_POST, onFirmwareUpdate, asyncWebuiUpdateHandler);
+
   server.on("/api/restart", HTTP_GET, onApiRestart);
   server.addRewrite(new OneParamRewrite("/api/lights/color/{color}",
                                         "/api/lights/color?c={color}"));
@@ -109,6 +112,89 @@ void setupWebserver()
 
 void stopWebServer() { server.end(); }
 
+void onFirmwareUpdate(AsyncWebServerRequest *request)
+{
+  bool shouldReboot = !Update.hasError();
+  AsyncWebServerResponse *response = request->beginResponse(200, "text/plain", shouldReboot ? "OK" : "FAIL");
+  response->addHeader("Connection", "close");
+  request->send(response);
+}
+
+
+
+void asyncWebuiUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final)
+{
+  asyncFileUpdateHandler(request, filename, index, data, len, final, U_SPIFFS);
+}
+
+void asyncFileUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final, int command)
+{
+    if (!index)
+  {
+    Serial.printf("Update Start: %s\n", filename.c_str());
+
+    // Update.runAsync(true);
+    if (!Update.begin((ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000), command)
+    {
+      Update.printError(Serial);
+    }
+  }
+  if (!Update.hasError())
+  {
+    if (Update.write(data, len) != len)
+    {
+      Update.printError(Serial);
+    }
+  }
+  if (final)
+  {
+    if (Update.end(true))
+    {
+      Serial.printf("Update Success: %uB\n", index + len);
+      onApiRestart(request);
+    }
+    else
+    {
+      Update.printError(Serial);
+    }
+  }
+}
+
+void asyncFirmwareUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final)
+{
+    asyncFileUpdateHandler(request, filename, index, data, len, final, U_FLASH);
+
+  // if (!index)
+  // {
+  //   Serial.printf("Update Start: %s\n", filename.c_str());
+
+  //   // Update.runAsync(true);
+  //   if (!Update.begin((ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000))
+  //   {
+  //     Update.printError(Serial);
+  //   }
+  // }
+  // if (!Update.hasError())
+  // {
+  //   if (Update.write(data, len) != len)
+  //   {
+  //     Update.printError(Serial);
+  //   }
+  // }
+  // if (final)
+  // {
+  //   if (Update.end(true))
+  //   {
+  //     Serial.printf("Update Success: %uB\n", index + len);
+  //     onApiRestart(request);
+  //   }
+  //   else
+  //   {
+  //     Update.printError(Serial);
+  //   }
+  // }
+}
+
 JsonDocument getStatusObject()
 {
   JsonDocument root;
diff --git a/src/lib/webserver.hpp b/src/lib/webserver.hpp
index 817869a..3d16b35 100644
--- a/src/lib/webserver.hpp
+++ b/src/lib/webserver.hpp
@@ -44,6 +44,10 @@ void onApiLightsSetColor(AsyncWebServerRequest *request);
 void onApiLightsSetJson(AsyncWebServerRequest *request, JsonVariant &json);
 
 void onApiRestart(AsyncWebServerRequest *request);
+void onFirmwareUpdate(AsyncWebServerRequest *request);
+void asyncFirmwareUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final);
+void asyncFileUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final, int command);
+void asyncWebuiUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final);
 
 void onIndex(AsyncWebServerRequest *request);
 void onNotFound(AsyncWebServerRequest *request);