#include <WiFi.h>
#include <ThingsBoard.h>
#include <Arduino_MQTT_Client.h>
#include <Espressif_Updater.h>

#define SERVER "demo.thingsboard.io"
#define TOKEN "445sb5cwshuymk3mdsd3"

constexpr uint16_t MAX_MESSAGE_SIZE = 512U;
constexpr char CURRENT_FIRMWARE_TITLE[] = "SinglePole";
constexpr char CURRENT_FIRMWARE_VERSION[] = "1.0";
constexpr uint8_t FIRMWARE_FAILURE_RETRIES = 12U;
constexpr uint16_t FIRMWARE_PACKET_SIZE = 4096U;

WiFiClient wifiClient;
Arduino_MQTT_Client mqttClient(wifiClient);
ThingsBoard tb(mqttClient, MAX_MESSAGE_SIZE);

Espressif_Updater updater;
const OTA_Update_Callback otaCallback(&progressCallback, &updatedCallback, CURRENT_FIRMWARE_TITLE, CURRENT_FIRMWARE_VERSION, &updater, FIRMWARE_FAILURE_RETRIES, FIRMWARE_PACKET_SIZE);

bool firmwareInfoSent = false;
bool updateRequestSent = false;

void progressCallback(const size_t& currentChunk, const size_t& totalChunks) {
    Serial.printf("Progress %.2f%%\n", static_cast<float>(currentChunk * 100U) / totalChunks);
}

void updatedCallback(const bool& success) {
    if (success) {
        Serial.println("Done, Reboot now");
        esp_restart();
        return;
    }
    Serial.println("Downloading firmware failed");
}

void connectToThingsBoard() {
  if (!tb.connected()) {
    Serial.println("Connecting to ThingsBoard server");
    if (!tb.connect(SERVER, TOKEN)) {
      Serial.println("Failed to connect to ThingsBoard");
    } else {
      Serial.println("Connected to ThingsBoard");
    }
  }
}

void handleOTAUpdate() {
    if (!firmwareInfoSent) {
        firmwareInfoSent = tb.Firmware_Send_Info(CURRENT_FIRMWARE_TITLE, CURRENT_FIRMWARE_VERSION) 
                        && tb.Firmware_Send_State(FW_STATE_UPDATING);
    }
    if (!updateRequestSent) {
        Serial.printf("Firmware Updated, version %s...", CURRENT_FIRMWARE_VERSION);
        updateRequestSent = tb.Subscribe_Firmware_Update(otaCallback);
    }
}

void setup() {
    WiFi.mode(WIFI_STA);
    Serial.begin(115200);
}

void loop() {
  if (!tb.connected()) {
    connectToThingsBoard();
  }

  handleOTAUpdate();
  tb.loop();
  delay(1000);
}