From 266528587aaef80186a41a45e9d322bfc050e42d Mon Sep 17 00:00:00 2001
From: Brett Christensen <brett.christensen@rea-group.com>
Date: Fri, 2 Feb 2024 10:27:50 +1100
Subject: [PATCH] add wifi monitoring loop

---
 Software/Software.ino                         |  1 +
 Software/src/devboard/webserver/webserver.cpp | 78 +++++++++++++------
 Software/src/devboard/webserver/webserver.h   |  9 +++
 3 files changed, 63 insertions(+), 25 deletions(-)

diff --git a/Software/Software.ino b/Software/Software.ino
index 63828a420..2c548ce69 100644
--- a/Software/Software.ino
+++ b/Software/Software.ino
@@ -148,6 +148,7 @@ void loop() {
 #ifdef WEBSERVER
   // Over-the-air updates by ElegantOTA
   ElegantOTA.loop();
+  WiFi_monitor_loop();
 #endif
 
   // Input
diff --git a/Software/src/devboard/webserver/webserver.cpp b/Software/src/devboard/webserver/webserver.cpp
index 21fa1afcc..d2bc3460e 100644
--- a/Software/src/devboard/webserver/webserver.cpp
+++ b/Software/src/devboard/webserver/webserver.cpp
@@ -42,7 +42,9 @@ bool wifi_connected;
 // Wifi connect time declarations and definition
 unsigned long wifi_connect_start_time;
 unsigned long wifi_connect_current_time;
-const long wifi_connect_timeout = 5000;  // Timeout for WiFi connect in milliseconds
+unsigned long wifi_connect_timeout = 5000;     // Timeout for WiFi connect in milliseconds
+unsigned long wifi_monitor_loop_time = 30000;  // Will check if WiFi is connected and try reconnect every x milliseconds
+unsigned long last_wifi_monitor_run = 0;
 
 void init_webserver() {
   // Configure WiFi
@@ -221,6 +223,23 @@ void init_webserver() {
   server.begin();
 }
 
+void WiFi_monitor_loop() {
+  unsigned long currentMillis = millis();
+  if (currentMillis - last_wifi_monitor_run > wifi_monitor_loop_time) {
+    last_wifi_monitor_run = currentMillis;
+    if (WiFi.status() != WL_CONNECTED && wifi_state != "Connecting") {
+      wifi_connected = false;
+      wifi_state = "Not connected";
+      Serial.print("Wifi disconnected. Attempting reconnection");
+      init_WiFi_STA(ssid, password);
+    } else if (WiFi.status() == WL_CONNECTED && wifi_state != "Connected") {
+      wifi_connected = true;
+      wifi_state = "Connected";
+      Serial.println("Wifi reconnected");
+    }
+  }
+}
+
 void init_WiFi_AP() {
   Serial.print("Creating Access Point: ");
   Serial.println(ssidAP);
@@ -236,36 +255,45 @@ void init_WiFi_AP() {
 }
 
 void init_WiFi_STA(const char* ssid, const char* password) {
-  // Connect to Wi-Fi network with SSID and password
-  Serial.print("Connecting to ");
-  Serial.println(ssid);
-  WiFi.begin(ssid, password);
-
-  wifi_connect_start_time = millis();
-  wifi_connect_current_time = wifi_connect_start_time;
-  while ((wifi_connect_current_time - wifi_connect_start_time) <= wifi_connect_timeout &&
-         WiFi.status() != WL_CONNECTED) {  // do this loop for up to 5000ms
-    // to break the loop when the connection is not established (wrong ssid or password).
-    delay(500);
-    Serial.print(".");
-    wifi_connect_current_time = millis();
+  // If we're already connected, there's nothing to do
+  if (WiFi.status() == WL_CONNECTED) {
+    if (wifi_state != "Connected") {
+      wifi_connected = true;
+      wifi_state = "Connected";
+      // Print local IP address and start web server
+      Serial.println("");
+      Serial.print("Connected to WiFi network: ");
+      Serial.println(ssid);
+      Serial.print("IP address: ");
+      Serial.println(WiFi.localIP());
+    }
+    return;
   }
-  if (WiFi.status() == WL_CONNECTED) {  // WL_CONNECTED is assigned when connected to a WiFi network
-    wifi_connected = true;
-    wifi_state = "Connected";
-    // Print local IP address and start web server
-    Serial.println("");
-    Serial.print("Connected to WiFi network: ");
+
+  // If we're not currently trying to connect, start the connection process
+  if (wifi_state != "Connecting") {
+    Serial.print("Connecting to: ");
     Serial.println(ssid);
-    Serial.print("IP address: ");
-    Serial.println(WiFi.localIP());
-  } else {
-    wifi_connected = false;
+    WiFi.begin(ssid, password);
+    wifi_state = "Connecting";
+    wifi_connect_start_time = millis();
+    return;
+  }
+
+  // If we've been trying to connect for more than 5000ms, give up
+  if (millis() - wifi_connect_start_time > wifi_connect_timeout) {
     wifi_state = "Not connected";
-    Serial.print("Not connected to WiFi network: ");
+    Serial.print("Failed to connect to WiFi network: ");
     Serial.println(ssid);
     Serial.println("Please check WiFi network name and password, and if WiFi network is available.");
+    Serial.print("Will try again in ");
+    Serial.print((wifi_monitor_loop_time - (millis() - last_wifi_monitor_run)) / 1000);
+    Serial.println(" seconds.");
+    return;
   }
+
+  // Otherwise, just print a dot to indicate that we're still trying to connect
+  Serial.print(".");
 }
 
 void init_ElegantOTA() {
diff --git a/Software/src/devboard/webserver/webserver.h b/Software/src/devboard/webserver/webserver.h
index 1cfda30c0..03c275c95 100644
--- a/Software/src/devboard/webserver/webserver.h
+++ b/Software/src/devboard/webserver/webserver.h
@@ -71,6 +71,15 @@ void init_WiFi_AP();
  */
 void init_WiFi_STA(const char* ssid, const char* password);
 
+/**
+ * @brief Monitoring loop for WiFi. Will attempt to reconnect to access point if the connection goes down.
+ *
+ * @param[in] void
+ * 
+ * @return void
+ */
+void WiFi_monitor_loop();
+
 /**
  * @brief Initialization function for ElegantOTA.
  *