/*******************************************************
 * ESP32 Azure IoT Hub MQTT Client using PubSubClient
 * 
 * Features:
 * - Connects to Wi-Fi
 * - Synchronizes time via NTP
 * - Generates SAS token for Azure IoT Hub
 * - Connects to Azure IoT Hub via MQTT over TLS
 * - Publishes D2C telemetry messages
 * - Subscribes to C2D messages and handles incoming messages
 *******************************************************/

#include <WiFi.h>
#include <WiFiClientSecure.h>
#include <PubSubClient.h>
#include <time.h>
#include <base64.h>
#include <mbedtls/md.h>

// Replace these with your actual credentials
#define WIFI_SSID           "Wokwi-GUEST"
#define WIFI_PASSWORD       ""

#define IOTHUB_HOSTNAME     "iot-devicetwin.azure-devices.net"
#define DEVICE_ID           "YOUR_DEVICE_ID"
#define DEVICE_KEY          "YOUR_DEVICE_PRIMARY_KEY"

#define MQTT_PORT           8883
#define TELEMETRY_INTERVAL  10000  // 10 seconds

// Azure IoT Hub MQTT Topics
#define D2C_TOPIC           "devices/" DEVICE_ID "/messages/events/"
#define C2D_SUBSCRIBE_TOPIC "devices/" DEVICE_ID "/messages/devicebound/#"

// Time Settings
#define TIME_ZONE_OFFSET    -8  // UTC offset in hours
#define DAYLIGHT_OFFSET     1   // Daylight offset in hours

// Function Prototypes
String generate_sas_token(const char* uri, const char* key, const int expiry);
String url_encode(String str);
String hmac_sha256(const char* key, const char* data);

// Global Variables
WiFiClientSecure secureClient;
PubSubClient mqttClient(secureClient);
unsigned long lastTelemetry = 0;

// Root CA certificate for Azure IoT Hub (Baltimore CyberTrust Root)
const char* root_ca = \
"-----BEGIN CERTIFICATE-----\n" \
"MIIDdzCCAl+gAwIBAgIEbOa7OzANBgkqhkiG9w0BAQsFADBvMQswCQYDVQQGEwJV\n" \
"UzELMAkGA1UECAwCTlkxEDAOBgNVBAcMB05ldyBZb3JrMRAwDgYDVQQKDAdNeUNv\n" \
"bXBhbnkxEDAOBgNVBAsMB015RGVwYXJ0MRowGAYDVQQDDBFteWNvbXBhbnkuY29t\n" \
"MB4XDTIxMDcxMzA5NTQ1M1oXDTMxMDcxMTA5NTQ1M1owbzELMAkGA1UEBhMCVVMx\n" \
"CzAJBgNVBAgMAk5ZMQwwCgYDVQQHDANOZXcxEDAOBgNVBAoMB015Q29tcGFueTEQ\n" \
"MA4GA1UECwwHTXlEZXBhcnQxGjAYBgNVBAMMEW15Y29tcGFueS5jb20wggEiMA0G\n" \
"CSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDnhN3T4s0A8kBYMwV/WYDyexH/t+Nn\n" \
"u+8J5GwR+yyuTcSfQk7iKTWyXj0nYrfNG5+1smWNGp29A6FWQwRZd/G6ZVYcXqAd\n" \
"qhj/TiVmqz5e8vV0FXXY/MWp8oOkOrjUymz6zPGqkPbXbzIYWSH6mCv0tJJ5zRlg\n" \
"biUAEaQjClQV6nH2FQhOW0V4cmEYXb6r1yb3u6ElpKUBG5/5u1JfKPWMSBNX2mVr\n" \
"Gf0J7JzvvTMJOscnBtTg/Mr6FExvzRTFzmwJzFqWQ2F0ODXgFg3GmG+otV/Nd2GE\n" \
"IaCqk5A1qfs7uE9VwV+/QnMtmV+9AoCTDT9iy5v5/BF8WKRtV1TnFh4BAgMBAAGj\n" \
"UDBOMB0GA1UdDgQWBBR+b6bMbldMg3Y/++C+RFsnAzSyGDAfBgNVHSMEGDAWgBR+\n" \
"b6bMbldMg3Y/++C+RFsnAzSyGDAMBgNVHRMEBTADAQH/MA0GCSqGSIb3DQEBCwUA\n" \
"A4IBAQA6T1WS8FD2dRRAkIg7k9aPH3FZ1q1FjxQgJZ8uVb76sPuImBICc3BQnFxJ\n" \
"vFo0c8gTfJdRXBx+O8Mlgqj1RHd7BkSXwVMGCt2dLklC03LwzVRshx1c3W8mFN8c\n" \
"yZtFzYAFdV5Ew/VzVtO0mkqk1DRSos5+ZmG7IsL5JqzQpQWbE+gD2hSfnRcNCM7S\n" \
"OH3JGh5FZgGcPQn0nE7QoX1UFm5wz3GJgZJYx1qJ/6eXHKUSOGpXn5BF/aIoHfhg\n" \
"eRMvDmzrj9ByJp1x3ZvxG5RT8vZg5hC5pBPCVn0NfVwQo9q1lqcXliCjBmQwP4UR\n" \
"V5AM7X6htlNCpGkxMldoZ1XxJ1BK\n" \
"-----END CERTIFICATE-----\n";

// Helper function to generate SAS token
String generate_sas_token(const char* uri, const char* key, const int expiry) {
  // Calculate expiry
  unsigned long currentTime = time(NULL);
  unsigned long expiryTime = currentTime + expiry;

  // Create string to sign
  String stringToSign = String(uri) + "\n" + String(expiryTime);

  // Decode the base64 device key
  int decodedLength = base64_decoded_length(String(key));
  uint8_t decodedKey[decodedLength];
  base64_decode(decodedKey, key, strlen(key));

  // HMAC-SHA256 signature
  uint8_t hash[32];  // SHA256 hash is 32 bytes
  mbedtls_md_context_t ctx;
  mbedtls_md_init(&ctx);
  mbedtls_md_setup(&ctx, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 1);
  mbedtls_md_hmac_starts(&ctx, decodedKey, decodedLength);
  mbedtls_md_hmac_update(&ctx, (const unsigned char*)stringToSign.c_str(), stringToSign.length());
  mbedtls_md_hmac_finish(&ctx, hash);
  mbedtls_md_free(&ctx);

  // Base64 encode the hash
  String encodedHash = base64::encode(hash, sizeof(hash));

  // URL encode the base64 string
  String encodedSig = url_encode(encodedHash);

  // Construct the SAS token
  String sasToken = "SharedAccessSignature sr=" + String(uri) + "&sig=" + encodedSig + "&se=" + String(expiryTime);
  return sasToken;
}

// Helper function to URL encode a string
String url_encode(String str) {
  String encoded = "";
  char c;
  char code0;
  char code1;
  for (int i = 0; i < str.length(); i++) {
    c = str.charAt(i);
    if (isalnum(c) || c == '-' || c == '_' || c == '.' || c == '~') {
      encoded += c;
    }
    else {
      code0 = (c & 0xf0) >> 4;
      code1 = (c & 0x0f);
      encoded += '%';
      encoded += (code0 < 10) ? ('0' + code0) : ('A' + code0 - 10);
      encoded += (code1 < 10) ? ('0' + code1) : ('A' + code1 - 10);
    }
  }
  return encoded;
}

void setup() {
  // Initialize Serial Monitor
  Serial.begin(115200);
  delay(1000);
  Serial.println();
  Serial.println("ESP32 Azure IoT Hub MQTT Client");

  // Connect to Wi-Fi
  Serial.print("Connecting to Wi-Fi: ");
  Serial.println(WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }

  Serial.println("\nWi-Fi connected!");
  Serial.print("IP Address: ");
  Serial.println(WiFi.localIP());

  // Initialize NTP
  configTime(TIME_ZONE_OFFSET * 3600, DAYLIGHT_OFFSET * 3600, "pool.ntp.org", "time.nist.gov");
  Serial.println("Synchronizing time with NTP...");
  struct tm timeinfo;
  if (!getLocalTime(&timeinfo)) {
    Serial.println("Failed to obtain time");
    return;
  }
  Serial.println("Time synchronized");
  Serial.print("Current Time: ");
  Serial.println(&timeinfo, "%Y-%m-%d %H:%M:%S");

  // Generate SAS Token
  String uri = String(IOTHUB_HOSTNAME) + "/devices/" + String(DEVICE_ID);
  String sas_token = generate_sas_token(uri.c_str(), DEVICE_KEY, SAS_TOKEN_TTL_SECS);
  Serial.println("SAS Token Generated:");
  Serial.println(sas_token);

  // Configure TLS
  secureClient.setCACert(root_ca);

  // Configure MQTT Client
  mqttClient.setServer(IOTHUB_HOSTNAME, MQTT_PORT);
  mqttClient.setCallback([](char* topic, byte* payload, unsigned int length) {
    // Handle incoming C2D messages
    Serial.print("Received message on topic: ");
    Serial.println(topic);
    String message;
    for (unsigned int i = 0; i < length; i++) {
      message += (char)payload[i];
    }
    Serial.print("Message payload: ");
    Serial.println(message);
    Serial.println("---------------");
  });

  // Connect to MQTT
  String mqtt_username = String(IOTHUB_HOSTNAME) + "/" + String(DEVICE_ID) + "/?api-version=2021-04-12";
  mqttClient.connect(DEVICE_ID, mqtt_username.c_str(), sas_token.c_str());

  if (mqttClient.connected()) {
    Serial.println("Connected to Azure IoT Hub via MQTT");
    // Subscribe to C2D messages
    mqttClient.subscribe(C2D_SUBSCRIBE_TOPIC);
    Serial.print("Subscribed to C2D topic: ");
    Serial.println(C2D_SUBSCRIBE_TOPIC);
  } else {
    Serial.println("Failed to connect to Azure IoT Hub");
  }
}

void loop() {
  // Maintain MQTT connection
  if (!mqttClient.connected()) {
    Serial.println("Disconnected from MQTT. Attempting to reconnect...");
    String uri = String(IOTHUB_HOSTNAME) + "/devices/" + String(DEVICE_ID);
    String sas_token = generate_sas_token(uri.c_str(), DEVICE_KEY, SAS_TOKEN_TTL_SECS);
    String mqtt_username = String(IOTHUB_HOSTNAME) + "/" + String(DEVICE_ID) + "/?api-version=2021-04-12";
    if (mqttClient.connect(DEVICE_ID, mqtt_username.c_str(), sas_token.c_str())) {
      Serial.println("Reconnected to Azure IoT Hub via MQTT");
      mqttClient.subscribe(C2D_SUBSCRIBE_TOPIC);
      Serial.print("Subscribed to C2D topic: ");
      Serial.println(C2D_SUBSCRIBE_TOPIC);
    } else {
      Serial.print("Failed to reconnect, rc=");
      Serial.println(mqttClient.state());
      delay(5000);
      return;
    }
  }

  mqttClient.loop();

  // Publish telemetry periodically
  unsigned long currentMillis = millis();
  if (currentMillis - lastTelemetry > TELEMETRY_INTERVAL) {
    lastTelemetry = currentMillis;

    // Example telemetry payload
    String telemetry = "{\"temperature\":25.3,\"humidity\":60}";

    if (mqttClient.publish(D2C_TOPIC, telemetry.c_str(), telemetry.length(), 1)) {
      Serial.print("Telemetry sent: ");
      Serial.println(telemetry);
    } else {
      Serial.println("Failed to send telemetry");
    }
  }
}