From bbf8e0b6b628bc053c67bfc29f2393555efb6ef5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kuba=20Szczodrzy=C5=84ski?= <kuba@szczodrzynski.pl>
Date: Mon, 6 Jun 2022 21:45:18 +0200
Subject: [PATCH] [realtek-ambz] Better handle WiFi events

---
 arduino/realtek-ambz/libraries/WiFi/WiFiAP.cpp      | 13 +++++++++++--
 arduino/realtek-ambz/libraries/WiFi/WiFiEvents.cpp  |  5 +----
 arduino/realtek-ambz/libraries/WiFi/WiFiGeneric.cpp | 12 ++++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/arduino/realtek-ambz/libraries/WiFi/WiFiAP.cpp b/arduino/realtek-ambz/libraries/WiFi/WiFiAP.cpp
index d055a8d..8efb06d 100644
--- a/arduino/realtek-ambz/libraries/WiFi/WiFiAP.cpp
+++ b/arduino/realtek-ambz/libraries/WiFi/WiFiAP.cpp
@@ -3,6 +3,11 @@
 #include "WiFi.h"
 #include "WiFiPriv.h"
 
+typedef struct {
+	int count;
+	rtw_mac_t mac_list[AP_STA_NUM];
+} client_info_t;
+
 bool WiFiClass::softAP(const char *ssid, const char *passphrase, int channel, bool ssidHidden, int maxClients) {
 	if (!enableAP(true))
 		return false;
@@ -61,6 +66,8 @@ bool WiFiClass::softAP(const char *ssid, const char *passphrase, int channel, bo
 		);
 	}
 
+	wifi_indication(WIFI_EVENT_CONNECT, NULL, ARDUINO_EVENT_WIFI_AP_START, -2);
+
 	if (ret < 0) {
 		LT_E("SoftAP failed; ret=%d", ret);
 		return false;
@@ -109,9 +116,11 @@ bool WiFiClass::softAPdisconnect(bool wifiOff) {
 }
 
 uint8_t WiFiClass::softAPgetStationNum() {
-	// TODO
 	// the struct is at wifi_conf.c:2576
-	return 0;
+	client_info_t info;
+	info.count = AP_STA_NUM;
+	wifi_get_associated_client_list(&info, sizeof(info));
+	return info.count;
 }
 
 IPAddress WiFiClass::softAPIP() {
diff --git a/arduino/realtek-ambz/libraries/WiFi/WiFiEvents.cpp b/arduino/realtek-ambz/libraries/WiFi/WiFiEvents.cpp
index 07a0c2a..5962ed1 100644
--- a/arduino/realtek-ambz/libraries/WiFi/WiFiEvents.cpp
+++ b/arduino/realtek-ambz/libraries/WiFi/WiFiEvents.cpp
@@ -145,10 +145,6 @@ void WiFiClass::handleRtwEvent(uint16_t event, char *data, int len, int flags) {
 	memset(&eventInfo, 0, sizeof(EventInfo));
 
 	switch (event) {
-		case WIFI_EVENT_CONNECT:
-			eventId = ARDUINO_EVENT_WIFI_STA_START;
-			break;
-
 		case WIFI_EVENT_DISCONNECT:
 		case WIFI_EVENT_RECONNECTION_FAIL:
 			eventId									 = ARDUINO_EVENT_WIFI_STA_DISCONNECTED;
@@ -189,6 +185,7 @@ void WiFiClass::handleRtwEvent(uint16_t event, char *data, int len, int flags) {
 			memcpy(eventInfo.wifi_ap_stadisconnected.mac, (const char *)data, 6);
 			break;
 
+		// case WIFI_EVENT_CONNECT:
 		// case WIFI_EVENT_SCAN_RESULT_REPORT:
 		// case WIFI_EVENT_SEND_ACTION_DONE:
 		// case WIFI_EVENT_RX_MGNT:
diff --git a/arduino/realtek-ambz/libraries/WiFi/WiFiGeneric.cpp b/arduino/realtek-ambz/libraries/WiFi/WiFiGeneric.cpp
index 10c76e9..1e6ae0d 100644
--- a/arduino/realtek-ambz/libraries/WiFi/WiFiGeneric.cpp
+++ b/arduino/realtek-ambz/libraries/WiFi/WiFiGeneric.cpp
@@ -45,6 +45,18 @@ bool WiFiClass::mode(WiFiMode mode) {
 		LT_E("Error while changing mode(%u)", mode);
 		return false;
 	}
+
+	// send STA start/stop events and AP stop event (start is handled in softAP())
+	if ((mode & WIFI_MODE_STA) != (currentMode & WIFI_MODE_STA)) {
+		if (mode & WIFI_MODE_STA)
+			wifi_indication(WIFI_EVENT_CONNECT, NULL, ARDUINO_EVENT_WIFI_STA_START, -2);
+		else
+			wifi_indication(WIFI_EVENT_CONNECT, NULL, ARDUINO_EVENT_WIFI_STA_STOP, -2);
+	}
+	if (!(mode & WIFI_MODE_AP) && (currentMode & WIFI_MODE_AP)) {
+		wifi_indication(WIFI_EVENT_CONNECT, NULL, ARDUINO_EVENT_WIFI_AP_STOP, -2);
+	}
+
 	LT_HEAP_I();
 	return true;
 }