From 268fd9986d73dd4772f69e853bdc8b230628a370 Mon Sep 17 00:00:00 2001
From: leblane <lewis@bootloader.io>
Date: Thu, 15 Jun 2023 19:44:47 +0300
Subject: [PATCH] Slightly more robust wifi initialisation. Also set country
 code.

---
 src/WatchFeatures/Wifi.cpp | 29 ++++++++++++++++++++++++++---
 src/WatchFeatures/Wifi.h   |  1 +
 src/Watchy.cpp             |  2 +-
 src/config.h               |  2 ++
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/WatchFeatures/Wifi.cpp b/src/WatchFeatures/Wifi.cpp
index 8eff1e0..d528572 100644
--- a/src/WatchFeatures/Wifi.cpp
+++ b/src/WatchFeatures/Wifi.cpp
@@ -2,14 +2,19 @@
 #include "Wifi.h"
 #include "config.h"
 
+#include <esp_err.h>
+#include <esp_wifi.h>
+
 bool WatchFeatures::Wifi::Connect()
 {
-  if (WiFi.begin(WIFI_SSID, WIFI_PASS) == WL_CONNECT_FAILED) {
+  WiFi.mode(WIFI_STA);
+  auto status = WiFi.begin(WIFI_SSID, WIFI_PASS);
+  if (status == WL_CONNECT_FAILED) {
     WiFi.mode(WIFI_OFF);
     return false;
   }
 
-  if (WiFi.waitForConnectResult() != WL_CONNECTED) {
+  if(WiFi.waitForConnectResult(5000) != WL_CONNECTED) {
     WiFi.mode(WIFI_OFF);
     return false;
   }
@@ -19,10 +24,28 @@ bool WatchFeatures::Wifi::Connect()
 
 void WatchFeatures::Wifi::Disconnect()
 {
-
+  WiFi.mode(WIFI_OFF);
 }
 
 bool WatchFeatures::Wifi::IsConnected()
 {
   return true;
+}
+
+void WatchFeatures::Wifi::InitBoot()
+{
+  WiFi.mode(WIFI_OFF);
+  WiFi.persistent(false);        // prevent the flash storage WiFi credentials
+  WiFi.setAutoReconnect(false);  // Set whether module will attempt to reconnect to an access point in case it is disconnected
+  WiFi.softAPdisconnect(false);  // kill rogue AP on startup
+  WiFi.disconnect(true);
+  WiFi.mode(WIFI_STA);
+
+  char country[3];
+  esp_wifi_get_country_code(country);
+  country[2] = '\0';
+
+  if (strcmp(country, WIFI_COUNTRY_CODE) != 0) {
+    esp_wifi_set_country_code(WIFI_COUNTRY_CODE, true);
+  }
 }
\ No newline at end of file
diff --git a/src/WatchFeatures/Wifi.h b/src/WatchFeatures/Wifi.h
index f173389..3c5f19f 100644
--- a/src/WatchFeatures/Wifi.h
+++ b/src/WatchFeatures/Wifi.h
@@ -8,5 +8,6 @@ namespace WatchFeatures
     bool Connect();
     void Disconnect();
     bool IsConnected();
+    void InitBoot();
   };
 }
\ No newline at end of file
diff --git a/src/Watchy.cpp b/src/Watchy.cpp
index a79fd78..20f6b72 100644
--- a/src/Watchy.cpp
+++ b/src/Watchy.cpp
@@ -134,9 +134,9 @@ void Watchy::InitBootInternal()
   m_display.epd2.selectSPI(SPI, SPISettings(20000000, MSBFIRST, SPI_MODE0));
   m_display.init(0, g_displayFullInit, 10, true);
   SetBusyCallback(); 
-
   BmaConfig();
   m_features.storage.InitBoot();
+  m_features.wifi.InitBoot();
   InitBoot();
 }
 
diff --git a/src/config.h b/src/config.h
index 0f29096..9340672 100644
--- a/src/config.h
+++ b/src/config.h
@@ -49,6 +49,8 @@
 #define DEFAULT_LONGITUDE 24.9375
 #define DEFAULT_CITY_NAME "Helsinki"
 
+#define WIFI_COUNTRY_CODE "FI"
+
 #include "secrets.h"
 #if !defined(WIFI_SSID) || !defined(WIFI_PASS)
 #error "Please define WIFI_SSID and WIFI_PASS in secrets.h"