Remote.h

// Copyright (C) 2024, Mark Qvist

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#include <WiFi.h>
#include <esp_wifi.h>
#if defined(UDP_TRANSPORT)
#include <WiFiUdp.h>
#include <microReticulum/Bytes.h>
#endif

#if CONFIG_IDF_TARGET_ESP32
  #include "esp32/rom/rtc.h"
#elif CONFIG_IDF_TARGET_ESP32S2
  #include "esp32s2/rom/rtc.h"
#elif CONFIG_IDF_TARGET_ESP32C3
  #include "esp32c3/rom/rtc.h"
#elif CONFIG_IDF_TARGET_ESP32S3
  #include "esp32s3/rom/rtc.h"
#else 
  #error Target CONFIG_IDF_TARGET is not supported
#endif

#define WIFI_UPDATE_INTERVAL_MS 500
#define WR_SOCKET_TIMEOUT 6
#define WR_READ_TIMEOUT_MS 6500
#define WR_RECONNECT_INTERVAL_MS 30000UL
#define WR_RECONNECT_MAX_INTERVAL_MS 120000UL
#define WR_STA_CONNECT_STUCK_MS 60000UL
#define WR_STA_RECOVERY_RESET_MS 300000UL

#ifndef WIFI_REMOTE_CONSOLE
  #define WIFI_REMOTE_CONSOLE 1
#endif

uint32_t wifi_update_interval_ms = WIFI_UPDATE_INTERVAL_MS;
uint32_t last_wifi_update = 0;
uint32_t wr_last_connect_try = 0;
uint32_t wr_last_read = 0;
uint32_t wr_last_disconnect_ms = 0;
uint32_t wr_last_got_ip_ms = 0;
uint32_t wr_sta_offline_since_ms = 0;
uint32_t wr_reconnect_interval_ms = WR_RECONNECT_INTERVAL_MS;
uint8_t wr_last_disconnect_reason = 0;
bool wr_have_disconnect_reason = false;
bool wr_have_got_ip = false;
bool wr_wifi_events_registered = false;

WiFiClient connection;
WiFiServer remote_listener(7633, 1);
IPAddress ap_ip(10, 0, 0, 1);
IPAddress ap_nm(255, 255, 255, 0);
IPAddress wr_device_ip;
char wr_hostname[10];
wl_status_t wr_wifi_status = WL_IDLE_STATUS;
#if defined(UDP_TRANSPORT)
WiFiUDP udp;
RNS::Bytes udp_buffer;
#if defined(HAS_RNS)
extern RNS::Interface udp_interface;
#endif
#endif

uint8_t wifi_mode = WIFI_OFF;
bool wifi_init_ran = false;
bool wifi_initialized = false;

char wr_ssid[33];
char wr_psk[33];

extern uint16_t udp_port;
extern void host_disconnected();

void wifi_dbg(String msg) { Serial.print("[WiFi] "); Serial.println(msg); }

uint8_t wifi_remote_mode() { return wifi_mode; }

bool wifi_is_connected() { return (wr_wifi_status == WL_CONNECTED); }
bool wifi_host_is_connected() {
#if WIFI_REMOTE_CONSOLE
  if (connection) { return true; } else { return false; }
#else
  return false;
#endif
}

void wifi_event_handler(WiFiEvent_t event, WiFiEventInfo_t info) {
  switch (event) {
    case ARDUINO_EVENT_WIFI_STA_CONNECTED:
      wr_have_got_ip = false;
      break;
    case ARDUINO_EVENT_WIFI_STA_GOT_IP:
      wr_last_got_ip_ms = millis();
      wr_sta_offline_since_ms = 0;
      wr_reconnect_interval_ms = WR_RECONNECT_INTERVAL_MS;
      wr_have_got_ip = true;
      wr_wifi_status = WL_CONNECTED;
      break;
    case ARDUINO_EVENT_WIFI_STA_DISCONNECTED:
      wr_last_disconnect_reason = info.wifi_sta_disconnected.reason;
      wr_last_disconnect_ms = millis();
      wr_have_disconnect_reason = true;
      wr_have_got_ip = false;
      wr_wifi_status = WiFi.status();
      break;
    default:
      break;
  }
}

void wifi_sta_note_reconnect_try(uint32_t now) {
  wr_last_connect_try = now;
  if (wr_reconnect_interval_ms < WR_RECONNECT_MAX_INTERVAL_MS) {
    wr_reconnect_interval_ms += wr_reconnect_interval_ms;
    if (wr_reconnect_interval_ms > WR_RECONNECT_MAX_INTERVAL_MS) {
      wr_reconnect_interval_ms = WR_RECONNECT_MAX_INTERVAL_MS;
    }
  }
}

void wifi_sta_soft_reconnect(uint32_t now) {
  if (wr_ssid[0] == 0x00) {
    return;
  }
  if ((uint32_t)(now - wr_last_connect_try) < wr_reconnect_interval_ms) {
    return;
  }
  wifi_sta_note_reconnect_try(now);
  // Runtime recovery must not tear the ESP32 WiFi driver down repeatedly.
  if (WiFi.getMode() & WIFI_MODE_STA) {
    if (WiFi.status() == WL_CONNECTED) {
      wr_reconnect_interval_ms = WR_RECONNECT_INTERVAL_MS;
      return;
    }
    bool sta_stuck = wr_sta_offline_since_ms != 0 && (uint32_t)(now - wr_sta_offline_since_ms) >= WR_STA_CONNECT_STUCK_MS;
    if (sta_stuck) {
      WARNING("WiFi STA connect stuck, restarting association");
      esp_wifi_disconnect();
      delay(200);
      if (wr_psk[0] != 0x00) { WiFi.begin(wr_ssid, wr_psk); }
      else                   { WiFi.begin(wr_ssid); }
      return;
    }
    esp_err_t err = esp_wifi_connect();
    if (err == ESP_OK || err == ESP_ERR_WIFI_CONN) {
      return;
    }
    WARNINGF("WiFi STA reconnect failed: 0x%x", err);
  } else {
    WiFi.mode(WIFI_STA);
    if (wr_psk[0] != 0x00) { WiFi.begin(wr_ssid, wr_psk); }
    else                   { WiFi.begin(wr_ssid); }
  }
}

void wifi_register_events() {
  if (!wr_wifi_events_registered) {
    WiFi.onEvent(wifi_event_handler);
    wr_wifi_events_registered = true;
  }
}

uint8_t wifi_last_disconnect_reason() {
  return wr_have_disconnect_reason ? wr_last_disconnect_reason : 0;
}

uint32_t wifi_last_disconnect_age_s() {
  if (!wr_have_disconnect_reason) return 0xFFFFFFFFUL;
  return (uint32_t)((millis() - wr_last_disconnect_ms) / 1000UL);
}

uint32_t wifi_connected_age_s() {
  if (!wr_have_got_ip || WiFi.status() != WL_CONNECTED) return 0xFFFFFFFFUL;
  return (uint32_t)((millis() - wr_last_got_ip_ms) / 1000UL);
}

int8_t wifi_rssi_dbm() {
  if (WiFi.status() != WL_CONNECTED) return 0;
  long rssi = WiFi.RSSI();
  if (rssi < -127) return -127;
  if (rssi > 127) return 127;
  return (int8_t)rssi;
}

void wifi_remote_start_ap() {
  WiFi.mode(WIFI_AP);
  if (wr_ssid[0] != 0x00) {
    if (wr_psk[0] != 0x00) { WiFi.softAP(wr_ssid, wr_psk, wr_channel); }
    else                   { WiFi.softAP(wr_ssid, NULL, wr_channel); }
  } else {
    if (wr_psk[0] != 0x00) { WiFi.softAP(bt_devname, wr_psk, wr_channel); }
    else                   { WiFi.softAP(bt_devname, NULL, wr_channel); }
  }
  delay(150);
  WiFi.softAPConfig(ap_ip, ap_ip, ap_nm);
  wifi_initialized = true;
}

void wifi_remote_start_sta() {
  wifi_register_events();
  wr_have_got_ip = false;
  wr_reconnect_interval_ms = WR_RECONNECT_INTERVAL_MS;
  #if HAS_BLUETOOTH || HAS_BLE == true
    if (bt_ready || bt_state != BT_STATE_OFF) { bt_stop(); }
  #endif
  WiFi.mode(WIFI_STA);
  WiFi.persistent(false);
  WiFi.setSleep(true);
  esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
  WiFi.setAutoReconnect(true);

  uint8_t ip[4]; bool ip_ok = true;
  for (uint8_t i = 0; i < 4; i++) { ip[i]  = EEPROM.read(config_addr(ADDR_CONF_IP+i)); }
  if (ip[0]==0x00 && ip[1]==0x00 && ip[2]==0x00 && ip[3]==0x00) { ip_ok = false; }
  if (ip[0]==0xFF && ip[1]==0xFF && ip[2]==0xFF && ip[3]==0xFF) { ip_ok = false; }

  uint8_t nm[4]; bool nm_ok = true;
  for (uint8_t i = 0; i < 4; i++) { nm[i]  = EEPROM.read(config_addr(ADDR_CONF_NM+i)); }
  if (nm[0]==0x00 && nm[1]==0x00 && nm[2]==0x00 && nm[3]==0x00) { nm_ok = false; }
  if (nm[0]==0xFF && nm[1]==0xFF && nm[2]==0xFF && nm[3]==0xFF) { nm_ok = false; }

  if (ip_ok && nm_ok) {
    IPAddress sta_ip(ip[0], ip[1], ip[2], ip[3]);
    IPAddress sta_nm(nm[0], nm[1], nm[2], nm[3]);
    WiFi.config(sta_ip, sta_ip, sta_nm);
  }

  WiFi.setMinSecurity(WIFI_AUTH_WPA2_PSK);
#if ESP_ARDUINO_VERSION >= ESP_ARDUINO_VERSION_VAL(3, 0, 0)
  WiFi.setPmf(true, false);
#endif

  //WiFi.onEvent([](WiFiEvent_t event, WiFiEventInfo_t info) {
  //  Serial.printf("[WiFi] event=%d\n", event);
  //});

  delay(100);
  Serial.print("[WiFi] ssid: ");
  Serial.println(wr_ssid);
  //Serial.print("[WiFi] psk: ");
  //Serial.println(wr_psk);
  if (wr_ssid[0] != 0x00) {
    if (wr_psk[0] != 0x00) { WiFi.begin(wr_ssid, wr_psk); }
    else                   { WiFi.begin(wr_ssid); }
  }
  
  delay(500);
  //delay(10000);
  wr_wifi_status = WiFi.status(); 
  if (wr_wifi_status == WL_CONNECTED) {
    wr_sta_offline_since_ms = 0;
    wr_have_got_ip = true;
    wr_last_got_ip_ms = millis();
  } else if (wr_sta_offline_since_ms == 0) {
    wr_sta_offline_since_ms = millis();
  }
  //Serial.print("[WiFi] status: ");
  //Serial.println(wr_wifi_status);
	//Serial.print("[WiFi] ip: ");
	//Serial.println(WiFi.localIP());
  wifi_initialized = true;
  wr_last_connect_try = millis();
}

void wifi_remote_stop() {
  WiFi.softAPdisconnect(true);
  WiFi.disconnect(true, true);
  WiFi.mode(WIFI_MODE_NULL);
  wifi_initialized = false;
}

void wifi_remote_start() {
  if      (wifi_mode == WR_WIFI_AP)  { wifi_remote_start_ap(); }
  else if (wifi_mode == WR_WIFI_STA) { wifi_remote_start_sta(); }
  else                               { wifi_remote_stop(); }

  if (wifi_initialized == true) {
#if WIFI_REMOTE_CONSOLE
    remote_listener.begin();
    remote_listener.setTimeout(WR_SOCKET_TIMEOUT);
    wr_state = WR_STATE_ON;
#else
    remote_listener.end();
    wr_state = WR_STATE_OFF;
#endif
#if defined(UDP_TRANSPORT)
    udp.begin(udp_port);
#endif
  } else {
    remote_listener.end();
    wr_state = WR_STATE_OFF;
#if defined(UDP_TRANSPORT)
    udp.stop();
#endif
  }
}

void wifi_remote_init() {
  //Serial.print("Initializing WiFi...\n");
  memcpy(wr_hostname, bt_devname, 5);
  memcpy(wr_hostname+5, bt_devname+6, 4);
  wr_hostname[9] = 0x00;
  WiFi.softAPdisconnect(true);
  WiFi.disconnect(true, true);
  WiFi.mode(WIFI_MODE_NULL);
  WiFi.setHostname(wr_hostname);

  wr_ssid[32] = 0x00; wr_psk[32] = 0x00;
  for (uint8_t i = 0; i < 32; i++) { wr_ssid[i] = EEPROM.read(config_addr(ADDR_CONF_SSID+i)); if (wr_ssid[i] == 0xFF) { wr_ssid[i] = 0x00; } }
  for (uint8_t i = 0; i < 32; i++) { wr_psk[i]  = EEPROM.read(config_addr(ADDR_CONF_PSK+i));  if (wr_psk[i]  == 0xFF) { wr_psk[i]  = 0x00; } }
  wr_channel = EEPROM.read(eeprom_addr(ADDR_CONF_WCHN)); if (wr_channel < 1 || wr_channel > 14) { wr_channel = WR_CHANNEL_DEFAULT; }
  wifi_remote_start();
  wifi_init_ran = true;
}

void wifi_remote_close_all() {
#if WIFI_REMOTE_CONSOLE
  // wifi_dbg("Close all"); // TODO: Remove debug
  if (connection) { connection.stop(); }
  WiFiClient client = remote_listener.available();
  while (client) { client.stop(); client = remote_listener.available(); }
  wr_state = WR_STATE_ON;
#endif
}

void wifi_remote_check_active() {
#if WIFI_REMOTE_CONSOLE
  if (millis()-wr_last_read >= WR_READ_TIMEOUT_MS) {
    // wifi_dbg("Connection activity timed out"); // TODO: Remove debug
    if (connection && connection.connected()) {
      connection.stop();
      wifi_remote_close_all();
      host_disconnected();
    }
  }
#endif
}

bool wifi_remote_available() {
#if WIFI_REMOTE_CONSOLE
  if (connection) {
    if (connection.connected()) {
      if (connection.available()) { wr_last_read = millis(); return true; }
      else                        { wifi_remote_check_active(); return false; }
    } else {
      // wifi_dbg("Client disconnected"); // TODO: Remove debug
      wifi_remote_close_all();
      return false;
    }
  } else {
    WiFiClient client = remote_listener.available();
    if (!client) { return false; }
    else {
      // wifi_dbg("Client connected"); // TODO: Remove debug
      connection = client;
      wr_state = WR_STATE_CONNECTED;
      wr_last_read = millis();
      if (connection.available()) { return true; }
      else                        { return false; }
    }
  }
#else
  return false;
#endif
}

uint8_t wifi_remote_read() {
#if WIFI_REMOTE_CONSOLE
  if (connection && connection.available()) { return connection.read(); }
  else {
    // wifi_dbg("Error: No data to read from TCP socket"); // TODO: Remove debug
    if (connection) { wifi_remote_close_all(); }
    return 0xC0;
  }
#else
  return 0xC0;
#endif
}

void wifi_remote_write(uint8_t byte) {
#if WIFI_REMOTE_CONSOLE
  if (connection) { connection.write(byte); }
#endif
}

void wifi_update_status() {
  wr_wifi_status = WiFi.status();
  //Serial.print("[WiFi] status: ");
  //Serial.println(wr_wifi_status);
  if (wr_wifi_status == WL_CONNECTED) {
    wr_device_ip = WiFi.localIP();
    //Serial.print("[WiFi] ip: ");
    //Serial.println(WiFi.localIP());
  }
  if (wifi_mode == WR_WIFI_AP && wifi_initialized) { wr_device_ip = WiFi.softAPIP(); wr_wifi_status = WL_CONNECTED; }
  if (wifi_init_ran && wifi_mode == WR_WIFI_STA && wr_wifi_status != WL_CONNECTED) {
    uint32_t now = millis();
    if (wr_sta_offline_since_ms == 0) {
      wr_sta_offline_since_ms = now;
    }
    if ((uint32_t)(now - wr_sta_offline_since_ms) >= WR_STA_RECOVERY_RESET_MS) {
      hard_reset(FW_RESET_REASON_WIFI_RECOVERY);
    }
    wifi_sta_soft_reconnect(now);
  } else if (wifi_mode == WR_WIFI_STA && wr_wifi_status == WL_CONNECTED) {
    wr_sta_offline_since_ms = 0;
    wr_reconnect_interval_ms = WR_RECONNECT_INTERVAL_MS;
  }
}

void update_wifi() {
#if defined(UDP_TRANSPORT)
  if (wifi_initialized) {
    if (udp.parsePacket() > 0) {
      size_t len = udp.read(udp_buffer.writable(MTU), MTU);
    if (len > 0) {
        udp_buffer.resize(len);
#if defined(HAS_RNS)
        if (udp_interface) {
          udp_interface.handle_incoming(udp_buffer);
        }
#endif
      }
    }
  }
#endif
  if (millis()-last_wifi_update >= wifi_update_interval_ms) {
    wifi_update_status();
    last_wifi_update = millis();
  }
}