/*************************************************************************/
/*    Copyright (C) 2022 Wizardry and Steamworks - License: GNU GPLv3    */
/*************************************************************************/
 
// The AP to connect to via Wifi.
#define STA_SSID ""
// The AP Wifi password.
#define STA_PSK ""
// The MQTT broker to connect to.
#define MQTT_HOST ""
// The MQTT broker username.
#define MQTT_USERNAME ""
// The MQTT broker password.
#define MQTT_PASSWORD ""
// The MQTT broker port.
#define MQTT_PORT 1883
// The default MQTT client ID is "esp-CHIPID" where CHIPID is the ESP8266
// or ESP32 chip identifier.
#define MQTT_CLIENT_ID() String("esp-" + String(GET_CHIP_ID(), HEX))
// The authentication password to use for OTA updates.
#define OTA_PASSWORD ""
// The OTA port on which updates take place.
#define OTA_PORT 8266
// The default topic that the sketch subscribes to is "esp/CHIPID" where
// CHIPID is the ESP8266 or ESP32 chip identifier.
#define MQTT_TOPIC() String("esp/" + String(GET_CHIP_ID(), HEX))
 
// Platform specific defines.
#if defined(ARDUINO_ARCH_ESP8266)
#define GET_CHIP_ID() (ESP.getChipId())
#elif defined(ARDUINO_ARCH_ESP32)
#define GET_CHIP_ID() ((uint16_t)(ESP.getEfuseMac()>>32))
#endif
 
// Miscellaneous defines.
//#define CHIP_ID_HEX (String(GET_CHIP_ID()).c_str())
#define HOSTNAME() String("esp-" + String(GET_CHIP_ID(), HEX))
 
// Platform specific libraries.
#if defined(ARDUINO_ARCH_ESP8266)
#include <ESP8266WiFi.h>
#include <ESP8266mDNS.h>
#elif defined(ARDUINO_ARCH_ESP32)
#include <WiFi.h>
#include <ESPmDNS.h>
#endif
// General libraries.
#include <WiFiUdp.h>
#include <ArduinoOTA.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#if defined(ARDUINO_ARCH_ESP32)
#include <FS.h>
#include <SPIFFS.h>
#endif
 
WiFiClient espClient;
PubSubClient mqttClient(espClient);
 
// Define GPIO pins for supported architectures.
#if defined(ARDUINO_ARCH_ESP8266)
int PINS[] = { D0, D1, D2, D3, D4, D5, D6, D7, D8 };
int ANAL[] = { A0 };
#elif defined(ARDUINO_ARCH_ESP32)
int PINS[] = { 
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 
               11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
               21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
               31, 32, 33, 34, 35, 36, 37, 38, 39 
             };
int ANAL[] = { 
                0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 
                0, 12, 13, 14, 15, 0, 0, 0, 0, 0, 
                0, 0, 0, 0, 25, 26, 27, 0, 0, 0,
                0, 32, 33, 34, 35, 36, 0, 0, 39 
             };
#endif
 
String mqttSerialize(StaticJsonDocument<256> msg) {
  char output[256];
  serializeJson(msg, output, 256);
  return String(output);
}
 
void mqttCallback(char *topic, byte *payload, unsigned int length) {
  String msgTopic = String(topic);
  // payload is not null terminated and casting will not work
  char msgPayload[length + 1];
  snprintf(msgPayload, length + 1, "%s", payload);
  Serial.println("Message received on topic: " + String(topic) + " with payload: " + String(msgPayload));
 
  // Parse the payload sent to the MQTT topic as a JSON document.
  StaticJsonDocument<256> doc;
  Serial.println("Deserializing message....");
  DeserializationError error = deserializeJson(doc, msgPayload);
  if (error) {
    Serial.println("Failed to parse MQTT payload as JSON: " + String(error.c_str()));
    return;
  }
 
  // Do not process messages without an action key.
  if (!doc.containsKey("action")) {
    return;
  }
 
  String action = (const char *)doc["action"];
  if(action == "set") {
      String state = (const char *)doc["state"];
      const int pin = (const int)doc["pin"];
 
      Serial.println("Setting pin: " + String(pin) + " to state: " + String(state));
 
      pinMode(PINS[pin], OUTPUT);
 
      if (state == "on") {
        digitalWrite(PINS[pin], HIGH);
        int status = digitalRead(PINS[pin]);
        Serial.println("Pin " + String(pin) + " state is now: " + String(status));
        return;
      }
 
      digitalWrite(PINS[pin], LOW);
      int status = digitalRead(PINS[pin]);
      Serial.println("Pin " + String(pin) + " state is now: " + String(status));
      return;
  }
 
  if(action == "get") {
      const int pin = (const int)doc["pin"];
 
      Serial.println("Getting pin: " + String(pin) + " state.");
 
      int status = digitalRead(PINS[pin]);
      Serial.println("Pin " + String(pin) + " state is now: " + String(status));
 
      // Announce the action.
      StaticJsonDocument<256> msg;
      msg["pin"] = pin;
      switch(status) {
        case 1:
          msg["state"] = "on";
          break;
        case 0:
          msg["state"] = "off";
          break;
        default:
          msg["state"] = "unknown";
          break;
      }
 
      mqttClient.publish(MQTT_TOPIC().c_str(), mqttSerialize(msg).c_str());
      return;
  }

  if(action == "push") {
      const int pin = (const int)doc["pin"];
      int time = (const int)doc["delay"];
      if(time <= 0) {
        time = 75;
      }
 
      Serial.println("Pushing pin: " + String(pin) + " with intermediate delay: " + String(time));
 
      pinMode(PINS[pin], OUTPUT);
      digitalWrite(PINS[pin], HIGH);
      delay(time);
      digitalWrite(PINS[pin], LOW);

      return;
  }

  if(action == "measure") {
      const int pin = (const int)doc["pin"];
      
      Serial.println("Getting analog pin: " + String(pin) + " state.");
      int analogValue = analogRead(ANAL[pin]);

      Serial.println("Value of analog pin " + String(pin) + " is: " + String(analogValue));

      // Announce the analog value.
      StaticJsonDocument<256> msg;
      msg["pin"] = pin;
      msg["value"] = analogValue;
 
      mqttClient.publish(MQTT_TOPIC().c_str(), mqttSerialize(msg).c_str());
      return;
  }
}
 
bool mqttConnect() {
  Serial.println("Attempting to connect to MQTT broker: " + String(MQTT_HOST));
  mqttClient.setServer(MQTT_HOST, MQTT_PORT);
 
  StaticJsonDocument<256> msg;
  if (mqttClient.connect(MQTT_CLIENT_ID().c_str(), MQTT_USERNAME, MQTT_PASSWORD)) {
    Serial.println("Established connection with MQTT broker using client ID: " + MQTT_CLIENT_ID());
    mqttClient.setCallback(mqttCallback);
    msg["action"] = "connected";
    mqttClient.publish(MQTT_TOPIC().c_str(), mqttSerialize(msg).c_str());
    Serial.println("Attempting to subscribe to MQTT topic: " + MQTT_TOPIC());
    if (!mqttClient.subscribe(MQTT_TOPIC().c_str())) {
      Serial.println("Failed to subscribe to MQTT topic: " + MQTT_TOPIC());
      return false;
    }
    Serial.println("Subscribed to MQTT topic: " + MQTT_TOPIC());
    msg.clear();
    msg["action"] = "subscribed";
    mqttClient.publish(MQTT_TOPIC().c_str(), mqttSerialize(msg).c_str());
    return true;
  }
 
  Serial.println("Connection to MQTT broker failed with MQTT client state: " + String(mqttClient.state()));
 
  return false;
}
 
bool loopWifiConnected() {
  // Process OTA loop first since emergency OTA updates might be needed.
  ArduinoOTA.handle();
 
  // Process MQTT client loop.
  if (!mqttClient.connected()) {
    // If the connection to the MQTT broker has failed then sleep before carrying on.
    if (!mqttConnect()) {
      return false;
    }
  }
  mqttClient.loop();
 
  return true;
}
 
void setup() {
  Serial.begin(115200);
  Serial.println("Booted, setting up Wifi in 10s...");
  delay(10000);
 
  WiFi.mode(WIFI_STA);
#if defined(ARDUINO_ARCH_ESP8266)
  WiFi.hostname(HOSTNAME().c_str());
#elif defined(ARDUINO_ARCH_ESP32)
  WiFi.setHostname(HOSTNAME().c_str());
#endif
  WiFi.begin(STA_SSID, STA_PSK);
  while (WiFi.waitForConnectResult() != WL_CONNECTED) {
    Serial.println("Failed to connect to Wifi, rebooting in 5s...");
    delay(5000);
    ESP.restart();
  }
 
  Serial.print("Connected to Wifi: ");
  Serial.println(WiFi.localIP());
 
  Serial.println("Setting up OTA in 10s...");
  delay(10000);
 
  // Port defaults to 8266
  ArduinoOTA.setPort(OTA_PORT);
 
  // Hostname defaults to esp-[ChipID]
  ArduinoOTA.setHostname(HOSTNAME().c_str());
 
  // Set the OTA password
  ArduinoOTA.setPassword(OTA_PASSWORD);
 
  ArduinoOTA.onStart([]() {
    switch (ArduinoOTA.getCommand()) {
      case U_FLASH: // Sketch
        Serial.println("OTA start updating sketch.");
        break;
#if defined(ARDUINO_ARCH_ESP8266)
      case U_FS:
#elif defined(ARDUINO_ARCH_ESP32)
      case U_SPIFFS:
#endif
        Serial.println("OTA start updating filesystem.");
        SPIFFS.end();
        break;
      default:
        Serial.println("Unknown OTA update type.");
        break;
    }
  });
  ArduinoOTA.onEnd([]() {
    Serial.println("OTA update complete.");
    SPIFFS.begin();
#if defined(ARDUINO_ARCH_ESP8266)
    // For what it's worth, check the filesystem on ESP8266.
    SPIFFS.check();
#endif
    ESP.restart();
  });
  ArduinoOTA.onProgress([](unsigned int progress, unsigned int total) {
    Serial.printf("OTA update progress: %u%%\r", (progress / (total / 100)));
  });
  ArduinoOTA.onError([](ota_error_t error) {
    Serial.printf("OTA update error [%u]: ", error);
    switch (error) {
      case OTA_AUTH_ERROR:
        Serial.println("OTA authentication failed");
        break;
      case OTA_BEGIN_ERROR:
        Serial.println("OTA begin failed");
        break;
      case OTA_CONNECT_ERROR:
        Serial.println("OTA connect failed");
        break;
      case OTA_RECEIVE_ERROR:
        Serial.println("OTA receive failed");
        break;
      case OTA_END_ERROR:
        Serial.println("OTA end failed");
        break;
      default:
        Serial.println("Unknown OTA failure");
        break;
    }
    ESP.restart();
  });
  ArduinoOTA.begin();
 
  // Set up MQTT client.
  mqttClient.setServer(MQTT_HOST, MQTT_PORT);
  mqttClient.setCallback(mqttCallback);
 
  // Touchdown.
  Serial.println("Setup complete.");
}
 
void loop() {
  // Check the Wifi connection status.
  int wifiStatus = WiFi.status();
  switch (wifiStatus) {
    case WL_CONNECTED:
      if (!loopWifiConnected()) {
        delay(1000);
        break;
      }
      delay(1);
      break;
    case WL_NO_SHIELD:
      Serial.println("No Wifi shield present.");
      goto DEFAULT_CASE;
      break;
    case WL_NO_SSID_AVAIL:
      Serial.println("Configured SSID not found.");
      goto DEFAULT_CASE;
      break;
    // Temporary statuses indicating transitional states.
    case WL_IDLE_STATUS:
    case WL_SCAN_COMPLETED:
      delay(1000);
      break;
    // Fatal Wifi statuses trigger a delayed ESP restart.
    case WL_CONNECT_FAILED:
    case WL_CONNECTION_LOST:
    case WL_DISCONNECTED:
    default:
      Serial.println("Wifi connection failed with status: " + String(wifiStatus));
DEFAULT_CASE:
      delay(10000);
      ESP.restart();
      break;
  }
}