/* PersWiFiManagerExt
   v1.0.0
   Based on PersWiFiManager v5.0.0 from https://r-downing.github.io/PersWiFiManager/
   Extended by Flo Kra
*/

#include "PersWiFiManagerExt.h"

#if defined(ESP32)
#include <esp_wifi.h>
#endif

#ifdef WIFI_HTM_PROGMEM
const char wifi_htm[] PROGMEM = R"=====(
<!DOCTYPE html>
<html>
  <head>
    <meta name="viewport" content="width=device-width, initial-scale=1, user-scalable=no"/>
    <title>ESP WiFi</title>
    <script>
      function g(i){
        return document.getElementById(i);
      };
      function p(t,l){
        if(confirm(t)) window.location=l;
      };
      function E(s){
        return document.createElement(s)
      };
      var S="setAttribute",A="appendChild",H="innerHTML",X,wl;
      function scan(){
        if(X) return;
        X=new XMLHttpRequest(),wl=document.getElementById('wl');
        wl[H]="Scanning...";
        X.onreadystatechange=function(){
          if (this.readyState==4&&this.status==200){
            X=0;wl[H]="";
            this.responseText.split("\n").forEach(function (e){
              let t=e.split(","), s=t.slice(2).join(',');
              var d=E('div'),i=E('a'),c=E('a');
              i[S]('class','s');
              c[S]('class','q');
              i.onclick=function(){
                g('s').value=s;
                g('p').focus();
              };
              i[A](document.createTextNode(s));
              c[H]=t[0]+"%"+(parseInt(t[1])?"\uD83D\uDD12":"\u26A0");
              wl[A](i); 
              wl[A](c);
              wl[A](document.createElement('br'));
            });
          }
        };
        X.open("GET","wifi/list",true);
        X.send();
      };
    </script>
    <style>
      input{
        padding:5px;font-size:1em;width:95%;
      }
      body{
        text-align:center;font-family:verdana;background-color:black;color:white;
      }
      a{
        color:#1fa3ec;
      }
      button{
        border:0;border-radius:0.3em;background-color:#1fa3ec;color:#fff;
        line-height:2.4em;font-size:1.2em;width:100%;display:block;
      }
      .q{
        float:right;
      }
      .s{
        display:inline-block;width:14em;overflow:hidden;text-overflow:ellipsis;white-space:nowrap;
      }
      #wl{
        line-height:1.5em;
      }
    </style>
  </head>
  <body>
    <div style='text-align:left;display:inline-block;width:320px;padding:5px'>
      <button onclick="scan()">&#x21bb; Scan</button>
      <p id='wl'></p>
      <form method='post' action='/wifi/connect'>
        <input id='s' name='n' length=32 placeholder='SSID'>
        <br>
        <input id='p' name='p' length=64 type='password' placeholder='password'>
        <br><br>
        <button type='submit'>Connect</button>
      </form>
      <br><br>
      <button onclick="p('Reboot device?','/wifi/rst')">Reboot</button>
      <br>
      <a href="javascript:history.back()">Back</a> |<a href="/">Home</a>
    </div>
  </body>
</html>
)=====";
#endif

#if defined(ESP8266)
PersWiFiManager::PersWiFiManager(ESP8266WebServer& s, DNSServer& d) {
#elif defined(ESP32)
PersWiFiManager::PersWiFiManager(WebServer& s, DNSServer& d) {
#endif
  _server = &s;
  _dnsServer = &d;
  _apPass = "";
  _freshConnectionAttempt = false;

  _w1Ssid = "";
  _w1Pass = "";
  _w2Ssid = "";
  _w2Pass = "";
  _httpUser = "";
  _httpPass = "";
  
  _connectNonBlock = true;
  _apStartTime = 0;
  _useSetWifi1Credentials = false;
  _useSetWifi2Credentials = false;

  _wifi1ConnectAttempts = 0;
  _wifi2ConnectAttempts = 0;
  _connectSetWifi = 0;
  _apModeFinished = false;

  _lastSuccessfulConnect = 0;
  _forceRetryWifi1LastTime = 0;

  _apTimeout = WIFI_AP_TIMEOUT;
  _apStartAfter = WIFI_CONNECT_TIMEOUT_STARTAP;
  _connCheckInterval = WIFI_CONNCHECK_INTERVAL;
  _rebootOnNoConnectAfter = WIFI_REBOOT_ON_NO_CONNECT_AFTER;
  _forceRetryWifi1Interval = WIFI_FORCE_RETRY_WIFI1_INTERVAL;

  _lastRunHandleWifi = 0;
  _lastConnectionCheckOK = 0;
  _wifiDisconnectedCounter = 0;

  _apModeEnabled = true;
  _httpAuth = false;
} //PersWiFiManager

bool PersWiFiManager::attemptConnection(const String& ssid, const String& pass) {
  //attempt to connect to wifi
  WiFi.mode(WIFI_STA);
  if (ssid.length()) {
    resetSettings(); // To avoid issues (experience from WiFiManager)
    if (pass.length()) WiFi.begin(ssid.c_str(), pass.c_str());
    else WiFi.begin(ssid.c_str());
    _connectStartTime = millis();
  } else {
    if((getSsid() == "") && (WiFi.status() != WL_CONNECTED)) { // No saved credentials from last connection, so skip trying to connect
      if(!_useSetWifi1Credentials && !_useSetWifi2Credentials) { 
        _connectStartTime = millis();
        _freshConnectionAttempt = true;
      #ifdef DEBUG
        Serial.println("WiFi: no saved credentials");
      #endif
        return false;
      }
    } else {
      if(!_useSetWifi1Credentials && !_useSetWifi2Credentials) { 
        // if no credentials were set using setWifi1Credentials/setWifi2Credentials
        // use saved credentials from last successful connection
        WiFi.begin();
        _connectStartTime = millis();
      #ifdef DEBUG
        Serial.println("WiFi: try connect using saved credentials");
      #endif
      }
    }
  }

  // if Wifi credentials were set using setWifi[x]Credentials method - use them
  if (_useSetWifi1Credentials || _useSetWifi2Credentials) {
    if (!_w1Ssid.length() && !_w2Ssid.length()) {
    #ifdef DEBUG
      Serial.println("WiFi: ERROR - both SSIDs empty");
    #endif
      // both configured WiFis have an empty SSID - no WiFi configured, directly switch to AP mode 
      _freshConnectionAttempt = true;
      _connectStartTime = millis();
    }
    else {
      if(_useSetWifi1Credentials && (_wifi1ConnectAttempts < SETWIFI_CONNECT_TRIES)) {
        _wifi1ConnectAttempts++;
        _connectSetWifi = 1;
    #ifdef DEBUG
        Serial.print("WiFi: trying WiFi-1... attempt: "); Serial.println(_wifi1ConnectAttempts);
    #endif
          resetSettings();
          delay(10);
          if (_w1Pass.length()) {
            //Serial.println("WiFi.begin(_w1Ssid.c_str(), _w1Pass.c_str());");
            WiFi.begin(_w1Ssid.c_str(), _w1Pass.c_str());
          } 
          else {
            //Serial.println("WiFi.begin(_w1Ssid.c_str());");
            WiFi.begin(_w1Ssid.c_str());
          }
          _connectStartTime = millis();
      }
      else if(_useSetWifi2Credentials && (_wifi2ConnectAttempts < SETWIFI_CONNECT_TRIES)) {
        _wifi2ConnectAttempts++;
        _connectSetWifi = 2;
    #ifdef DEBUG
        Serial.print("WiFi: trying WiFi-2... attempt: "); Serial.println(_wifi2ConnectAttempts);
    #endif
        resetSettings();
        delay(10);
          if (_w2Pass.length()) {
            //Serial.println("WiFi.begin(_w2Ssid.c_str(), _w2Pass.c_str());");
            WiFi.begin(_w2Ssid.c_str(), _w2Pass.c_str());
          }
          else {
            //Serial.println("WiFi.begin(_w2Ssid.c_str());");
            WiFi.begin(_w2Ssid.c_str());
          }
          _connectStartTime = millis();
      }
    }
  }

  //if in nonblock mode, skip this loop
  while (!_connectNonBlock && _connectStartTime) {
    handleWiFi();
    delay(10);
  }

  return (WiFi.status() == WL_CONNECTED);

} //attemptConnection

void PersWiFiManager::handleWiFi() {
  if (_connectStartTime && WiFi.status() == WL_CONNECTED) {
    _connectStartTime = 0;
    _apModeFinished = true; // connection was successful - disable AP until reboot
    _lastSuccessfulConnect = millis();
#ifdef DEBUG
    if(_connectSetWifi > 0) Serial.print("WiFi: connected to WiFi-"); Serial.println(_connectSetWifi);
#endif
    if(_connectSetWifi == 1) _wifi1ConnectAttempts = 0;
    if(_connectSetWifi == 2) _wifi2ConnectAttempts = 0;
    if (_connectHandler) _connectHandler();
  }

  //if failed or not connected and time is up - try connecting again (only if currently using setWifi1/setWifi2, to possibly be successful with the 2nd config set)
  if(_useSetWifi1Credentials || _useSetWifi2Credentials ) {
    if ((_connectStartTime && (WiFi.status() == WL_CONNECT_FAILED)) || (_connectStartTime && (WiFi.status() != WL_CONNECTED) && ((millis() - _connectStartTime) > (1000 * WIFI_CONNECT_TIMEOUT)))) {
      attemptConnection();
    }
  }

  //start AP mode
  //if connect failed or no saved SSID or no WiFi credentials were found or not connected and time is up
  if(_apModeEnabled && !_apModeFinished && _connectStartTime) {
    if ((WiFi.status() == WL_CONNECT_FAILED) || _freshConnectionAttempt || ((WiFi.status() != WL_CONNECTED) && ((millis() - _connectStartTime) > (1000 * _apStartAfter)))) {
      startApMode();
      _connectStartTime = 0; //reset connect start time
      _apStartTime = millis();
      _freshConnectionAttempt = false;
    }
  }

  // if AP is on and AP timeout exceeded, switch off AP mode and start normal connection handling again
  if ( _apModeEnabled && _apTimeout > 0 && !_apModeFinished && _apStartTime && ((millis() - _apStartTime) > (60000 * _apTimeout))) {
    if (WiFi.softAPgetStationNum() > 0) { 
      // there are active connections - extend AP timeout
      _apStartTime = millis();
    }
    else { // no one is connected and timeout exceeded - stop AP
      stopApMode();
      _apModeFinished = true; // prevent AP again until reboot
      _apStartTime = 0;
    }
  }

  // run if WIFI_CONNCHECK_INTERVAL exceeded
  // try reconnect if not connected
  if( _connCheckInterval > 0 && (millis() - _lastRunHandleWifi) > (1000 * _connCheckInterval)) {
    _lastRunHandleWifi = millis();
    int currWifiState = WiFi.status();

#ifdef DEBUG_2
    Serial.print("WiFi: status: ");
    Serial.print(currWifiState);
    Serial.print("  _connectStartTime: ");
    Serial.print(_connectStartTime);
    Serial.print("  _lastSuccessfulConnect: ");
    Serial.print(_lastSuccessfulConnect);
    Serial.print("  _apStartTime: ");
    Serial.print(_apStartTime);
    Serial.print("  _apModeFinished: ");
    Serial.println(_apModeFinished);
#endif

#ifdef DEBUG
    if(currWifiState != _wifiLastState) {
      Serial.print("WiFi: status changed to: ");
      //Serial.println(currWifiState);
      if(currWifiState == WL_IDLE_STATUS) Serial.println(F("WL_IDLE_STATUS"));
      else if(currWifiState == WL_NO_SSID_AVAIL) Serial.println(F("WL_NO_SSID_AVAIL"));
      else if(currWifiState == WL_CONNECTED) Serial.println(F("WL_CONNECTED"));
      else if(currWifiState == WL_CONNECT_FAILED) Serial.println(F("WL_CONNECT_FAILED"));
      else if(currWifiState == WL_DISCONNECTED) Serial.println(F("WL_DISCONNECTED"));
    }
#endif

    if(currWifiState == WL_CONNECTED) _lastConnectionCheckOK = millis();

    // if enabled and connected to Wifi2 and _forceRetryWifi1Interval is exceeded
    // try reconnecting in order to prefer Wifi1 when set and available
    if(_forceRetryWifi1Interval && _apStartTime == 0) {
      if(_useSetWifi1Credentials && _connectSetWifi != 1 && _apModeFinished && _w1Ssid.length() && ((millis() - _lastSuccessfulConnect) > (60000 * _forceRetryWifi1Interval))) {
        if((millis() - _forceRetryWifi1LastTime) > 120000) {
        #ifdef DEBUG
          Serial.println("WiFi: connected to WiFi-2 - attempt reconnect WiFi-1");
        #endif
          _forceRetryWifi1LastTime = millis();
          _wifi1ConnectAttempts = 0;
          _wifi2ConnectAttempts = 0;
          attemptConnection();
        }
      }
    }

    // after AP mode is turned off, reconnect connection if it drops
    if(_apModeFinished && _apStartTime == 0 && currWifiState != WL_CONNECTED && _wifiLastState != WL_CONNECTED) {
      if( _wifiDisconnectedCounter > 2) {
      #ifdef DEBUG
        Serial.println("WiFi: not connected - attempt reconnect");
      #endif
        _wifi1ConnectAttempts = 0;
        _wifi2ConnectAttempts = 0;
        _wifiDisconnectedCounter = 0;
        attemptConnection();
      }
      else {
      #ifdef DEBUG_2
        Serial.print("WiFi: not connected - _wifiDisconnectedCounter=");
        Serial.print(_wifiDisconnectedCounter);
        Serial.println(" - _wifiDisconnectedCounter++");
      #endif
        _wifiDisconnectedCounter++;
      }
    }

    if(_rebootOnNoConnectAfter > 0 && ((millis() - _lastConnectionCheckOK) > (60000 * _rebootOnNoConnectAfter))) {
#ifdef DEBUG
      Serial.println("WiFi: could not connect for a long time - rebooting...");
#endif
      delay(100);
      //ESP.restart();
      // Adding Safer Restart method
#if defined(ESP8266)
	    ESP.wdtDisable();
	    ESP.reset();
#elif defined(ESP32)
      ESP.restart();
#endif
  	  delay(2000);
    }

    if(currWifiState != _wifiLastState) {
      _wifiLastState = currWifiState;
    }
  }
  
} //handleWiFi

void PersWiFiManager::startApMode(){
  //start AP mode
  
  IPAddress apIP(192, 168, 4, 1);
  WiFi.mode(WIFI_AP);
  WiFi.softAPConfig(apIP, apIP, IPAddress(255, 255, 255, 0));
  _apPass.length() ? WiFi.softAP(getApSsid().c_str(), _apPass.c_str()) : WiFi.softAP(getApSsid().c_str());
  #ifdef DEBUG
  Serial.print("WiFi: AP mode started, SSID: '");
  Serial.print(getApSsid().c_str());
  Serial.print("', PW: '");
  Serial.print(_apPass.c_str());
  Serial.println("'");
  #endif
  if (_apHandler) _apHandler();  
}//startApMode

void PersWiFiManager::stopApMode() {
  #ifdef DEBUG
  Serial.println("WiFi: AP mode stopped");
  #endif
  WiFi.softAPdisconnect(true);
  delay(10);
  attemptConnection(); // sets WiFi mode to STA and begins reconnecting
}

void PersWiFiManager::setConnectNonBlock(bool b) {
  _connectNonBlock = b;
} //setConnectNonBlock

void PersWiFiManager::setApModeAutostart(bool m) {
  _apModeEnabled = m;
}

void PersWiFiManager::setApTimeout(int t) {
  _apTimeout = t;
} //setApTimeout

void PersWiFiManager::setApStartAfter(int t) {
  _apStartAfter = t;
}

void PersWiFiManager::setConnCheckInterval(int t) {
  _connCheckInterval = t;
}

void PersWiFiManager::setForceRetryWifi1Interval(int t) {
  _forceRetryWifi1Interval = t;
}

void PersWiFiManager::setupWiFiHandlers() {
  IPAddress apIP(192, 168, 4, 1);
  _dnsServer->setErrorReplyCode(DNSReplyCode::NoError);
  _dnsServer->start((byte)53, "*", apIP); //used for captive portal in AP mode

  _server->on("/wifi/list", [&] () {
    bool requestOk = false;
    if(_httpAuth) {
      if (_server->authenticate(_httpUser.c_str(), _httpPass.c_str())) requestOk = true;
      else _server->send (401, "text/plain", "UNAUTHORIZED");
    }
    else requestOk = true;
    if(requestOk) {
      //scan for wifi networks
      int n = WiFi.scanNetworks();

      //build array of indices
      int ix[n];
      for (int i = 0; i < n; i++) ix[i] = i;

      //sort by signal strength
      for (int i = 0; i < n; i++) for (int j = 1; j < n - i; j++) if (WiFi.RSSI(ix[j]) > WiFi.RSSI(ix[j - 1])) std::swap(ix[j], ix[j - 1]);
      //remove duplicates
      for (int i = 0; i < n; i++) for (int j = i + 1; j < n; j++) if (WiFi.SSID(ix[i]).equals(WiFi.SSID(ix[j])) && WiFi.encryptionType(ix[i]) == WiFi.encryptionType(ix[j])) ix[j] = -1;

      //build plain text string of wifi info
      //format [signal%]:[encrypted 0 or 1]:SSID
      String s = "";
      s.reserve(2050);
      for (int i = 0; i < n && s.length() < 2000; i++) { //check s.length to limit memory usage
        if (ix[i] != -1) {
#if defined(ESP8266)
          s += String(i ? "\n" : "") + ((constrain(WiFi.RSSI(ix[i]), -100, -50) + 100) * 2) + ","
               + ((WiFi.encryptionType(ix[i]) == ENC_TYPE_NONE) ? 0 : 1) + "," + WiFi.SSID(ix[i]);
#elif defined(ESP32)
          s += String(i ? "\n" : "") + ((constrain(WiFi.RSSI(ix[i]), -100, -50) + 100) * 2) + ","
               + ((WiFi.encryptionType(ix[i]) == WIFI_AUTH_OPEN) ? 0 : 1) + "," + WiFi.SSID(ix[i]);
#endif
        }
      }
      //send string to client
      _server->send(200, "text/plain", s);
    }
  }); //_server->on /wifi/list

  _server->on("/wifi/connect", [&]() {
    bool requestOk = false;
    if(_httpAuth) {
      if (_server->authenticate(_httpUser.c_str(), _httpPass.c_str())) requestOk = true;
      else _server->send (401, "text/plain", "UNAUTHORIZED");
    }
    else requestOk = true;
    if(requestOk) {
      _server->send(200, "text/html", "connecting...");
      attemptConnection(_server->arg("n"), _server->arg("p"));
    }
  }); //_server->on /wifi/connect

  _server->on("/wifi/ap", [&](){
    bool requestOk = false;
    if(_httpAuth) {
      if (_server->authenticate(_httpUser.c_str(), _httpPass.c_str())) requestOk = true;
      else _server->send (401, "text/plain", "UNAUTHORIZED");
    }
    else requestOk = true;
    if(requestOk) {
      _server->send(200, "text/html", "access point: "+getApSsid());
      startApMode();
    }
  }); //_server->on /wifi/ap

  _server->on("/wifi/rst", [&]() {
    bool requestOk = false;
    if(_httpAuth) {
      if (_server->authenticate(_httpUser.c_str(), _httpPass.c_str())) requestOk = true;
      else _server->send (401, "text/plain", "UNAUTHORIZED");
    }
    else requestOk = true;
    if(requestOk) {
      _server->send(200, "text/html", "Rebooting...");
      delay(100);
      //ESP.restart();
      // Adding Safer Restart method
#if defined(ESP8266)
	    ESP.wdtDisable();
	    ESP.reset();
#elif defined(ESP32)
      ESP.restart();
#endif
	  delay(2000);
    }
  });

#ifdef WIFI_HTM_PROGMEM
  _server->on(WIFI_HTM_PATH, [&]() {
    bool requestOk = false;
    if(_httpAuth) {
      if (_server->authenticate(_httpUser.c_str(), _httpPass.c_str())) requestOk = true;
      else _server->requestAuthentication();
    }
    else requestOk = true;
    if(requestOk) {
      _server->sendHeader("Cache-Control", " no-cache, no-store, must-revalidate");
      _server->sendHeader("Expires", " 0");
      _server->send(200, "text/html", wifi_htm);
    }
  });
#endif

}//setupWiFiHandlers

bool PersWiFiManager::begin(const String& ssid, const String& pass) {
#if defined(ESP32)
    WiFi.mode(WIFI_STA);  // ESP32 needs this before setupWiFiHandlers(). Might be good for ESP8266 too?
#endif
  setupWiFiHandlers();
  return attemptConnection(ssid, pass); //switched order of these two for return
} //begin

void PersWiFiManager::resetSettings() {
#if defined(ESP8266)
  WiFi.disconnect();
#elif defined(ESP32)
  wifi_mode_t m = WiFi.getMode();
  if(!(m & WIFI_MODE_STA)) WiFi.mode(WIFI_STA);
  WiFi.disconnect(false, true);
  if(!(m & WIFI_MODE_STA)) WiFi.mode(m);
#endif
} // resetSettings

String PersWiFiManager::getApSsid() {
#if defined(ESP8266)
  return _apSsid.length() ? _apSsid : "ESP8266";
#elif defined(ESP32)
  return _apSsid.length() ? _apSsid : "ESP32";
#endif
} //getApSsid

String PersWiFiManager::getSsid() {
#if defined(ESP8266)
  return WiFi.SSID();
#elif defined(ESP32)
  wifi_config_t conf;
  esp_wifi_get_config(WIFI_IF_STA, &conf);  // load wifi settings to struct comf
  const char *SSID = reinterpret_cast<const char*>(conf.sta.ssid);
  return String(SSID);
#endif
} //getSsid

void PersWiFiManager::setApCredentials(const String& apSsid, const String& apPass) {
  if (apSsid.length()) _apSsid = apSsid;
  if (apPass.length() >= 8) _apPass = apPass;
} //setApCredentials

void PersWiFiManager::setWifi1Credentials(const String& wSsid, const String& wPass) {
  _useSetWifi1Credentials = true;
  //Serial.println("setWifi1Credentials -> _useSetWifi1Credentials=true");
  if (wSsid.length()) {
    _w1Ssid = wSsid;
  }
  if (wPass.length()) _w1Pass = wPass;
} //setWifi1Credentials

void PersWiFiManager::setWifi2Credentials(const String& wSsid, const String& wPass) {
  _useSetWifi2Credentials = true;
  //Serial.println("setWifi2Credentials -> _useSetWifi2Credentials=true");
  if (wSsid.length()) {
    _w2Ssid = wSsid;
  }
  if (wPass.length()) _w2Pass = wPass;
} //setWifi2Credentials

void PersWiFiManager::onConnect(WiFiChangeHandlerFunction fn) {
  _connectHandler = fn;
}

void PersWiFiManager::onAp(WiFiChangeHandlerFunction fn) {
  _apHandler = fn;
}

void PersWiFiManager::onApOff(WiFiChangeHandlerFunction fn) {
  _apOffHandler = fn;
}

void PersWiFiManager::setHttpCredentials(const String& user, const String& pass) {
  if(user.length() && pass.length()) {
    _httpUser = user;
    _httpPass = pass;
    _httpAuth = true;
  }
}

void PersWiFiManager::setRebootOnNoConnect(int t) {
  _rebootOnNoConnectAfter = t;
}

int PersWiFiManager::getActiveWiFiNum() {
  return _connectSetWifi;
}
