#define WIFI_SSID     "HomeAP"  // Accesspoint name
#define WIFI_PWD      "Secure"  // Password
#define WIFI_HOSTNAME "RROTA"   // ESP hostname
#define WIFI_WPS      true      // true or false for ESP8266 for using WPS

#ifdef ESP32
  #include <WiFi.h>
  #include <ESPmDNS.h>
#else
  #include <ESP8266WiFi.h>
  #include <ESP8266mDNS.h>
#endif

#include <WiFiUdp.h>
#include <ArduinoOTA.h>

// Global variables:
static bool didInitOTA = false;

bool startWiFi() {
  bool wpsSuccess = false;
  WiFi.begin(WIFI_SSID, WIFI_PWD);
  delay(1);
  #ifdef ESP32
  WiFi.setHostname(WIFI_HOSTNAME);
  #else
  WiFi.hostname(WIFI_HOSTNAME);
  #endif
  delay(1);
  int retries = 20;
  for( int i = 0; i < retries && WiFi.status() != WL_CONNECTED; i++ ) {
    if( i+1 < retries ) {
      Serial.printf("WiFi waiting to connect: %s...\n", WIFI_SSID);
      delay(500);
    }
  }
  if( WiFi.status() == WL_CONNECTED ) {
    Serial.printf("WiFi connected to %s\n", WIFI_SSID);
    wpsSuccess = true;
  }
  return wpsSuccess;
}

//======================================== 
// Scan for WPS connections 
//========================================
bool startWPS() {
  bool wpsSuccess = false;
  #ifdef ESP32
    wpsSuccess = startWiFi();
  #else
  if( WIFI_WPS ) {
    wpsSuccess = WiFi.beginWPSConfig();
    if( wpsSuccess ) {
      String newSSID = WiFi.SSID();
      if( newSSID.length() > 0 ) {
        Serial.printf("WPS ready. successfully connected to SSID [%s]\n", newSSID.c_str() );
        wpsSuccess = true;
      } 
      else {
        wpsSuccess = false;
      }
    }
  } 
  else {
    wpsSuccess = startWiFi();
  }
  #endif
  return wpsSuccess; 
}


//======================================== 
// Sketch Setup 
//========================================
void setup() {
  Serial.begin(115200);
  while( !Serial ); 
  Serial.setDebugOutput(true); // Makes the WPS scanning visible...
  
  WiFi.mode(WIFI_STA);

  // Scan for WPS connections:
  while( !startWPS() ) {
    delay(1000);
  }
}


//======================================== 
// Sketch Loop 
//========================================
void loop() {
  // Wait for an IP address:
  if( WiFi.status() != WL_CONNECTED ) {
    delay(100);
    return;
  }

  // Init OTA:
  if( !didInitOTA ) {
    didInitOTA = true;
    Serial.printf("setup OTA...%s\n", WiFi.localIP().toString().c_str() );  

    ArduinoOTA.onStart([]() {
      String type;
      if (ArduinoOTA.getCommand() == U_FLASH) {
        type = "sketch";
      } 
      else { // U_FS
        type = "filesystem";
      }
      // NOTE: if updating FS this would be the place to unmount FS using FS.end()
      Serial.println("Start updating " + type);
    });

    ArduinoOTA.onEnd([]() {
      Serial.println("\nEnd");
    });

    ArduinoOTA.onProgress([](unsigned int progress, unsigned int total) {
      Serial.printf("Progress: %u%%\n", (progress / (total / 100)));
    });

    ArduinoOTA.onError([](ota_error_t error) {
      Serial.printf("Error[%u]: ", error);
      if (error == OTA_AUTH_ERROR) {
        Serial.println("Auth Failed");
      } 
      else if (error == OTA_BEGIN_ERROR) {
        Serial.println("Begin Failed");
      } 
      else if (error == OTA_CONNECT_ERROR) {
        Serial.println("Connect Failed");
      } 
      else if (error == OTA_RECEIVE_ERROR) {
        Serial.println("Receive Failed");
      } 
      else if (error == OTA_END_ERROR) {
        Serial.println("End Failed");
      }
    });

    delay(100);
    ArduinoOTA.setHostname("RROTA");
    ArduinoOTA.begin();
    return;
  }

  // Handle the OTA packets:
  ArduinoOTA.handle();
}