/*
  Azure Event Grid MQTT Endpoint Probe with mTLS
  - Wi-Fi connect
  - SNTP time sync
  - DNS resolve
  - TCP reachability (port 8883)
  - TLS (server-only) handshake using CRT bundle (or custom CA)
  - TLS (mTLS) handshake with client certificate & private key

  Notes:
    - Client certificate must be PEM and match private key.
    - Private key must be PEM and UNENCRYPTED (no passphrase).
    - SNI uses HOSTNAME automatically; do NOT use raw IP.
*/
#include <Arduino.h>
#include <WiFi.h>
#include <WiFiClient.h>
#include <WiFiClientSecure.h>

#include <../constants.h>
#include <../secrets.h>

extern "C" {
  #include <lwip/netdb.h>
  #include <lwip/sockets.h>
  #include <lwip/inet.h>
  #include <lwip/errno.h>
  #include <time.h>
}
static const char* HOSTNAME  = "ThisIsNotTheMQTTBrokerYouAreLookingFor.newzealandnorth-1.ts.eventgrid.azure.net";
static const uint16_t PORT   = 8883;

// Time servers (for TLS validity window)
static const char* NTP_1 = "pool.ntp.org";
static const char* NTP_2 = "time.cloudflare.com";

static const char* errnoName(int e) {
  switch (e) {
    case 5:   return "EIO";
    case 101: return "ENETUNREACH";
    case 104: return "ECONNRESET";
    case 110: return "ETIMEDOUT";
    case 111: return "ECONNREFUSED";
    case 113: return "EHOSTUNREACH";
    default:  return "?";
  }
}


bool waitForWifi(uint32_t timeout_ms = 20000) {
  uint32_t start = millis();
  Serial.printf("[WiFi] Connecting to '%s'...\n", WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  while (WiFi.status() != WL_CONNECTED && (millis() - start) < timeout_ms) {
    delay(250);
    Serial.print(".");
  }
  Serial.println();
  return WiFi.status() == WL_CONNECTED;
}


void syncTime() {
  configTime(0, 0, NTP_1, NTP_2);
  Serial.println("[NTP] Syncing time...");
  for (int i = 0; i < 20; ++i) {
    time_t now = time(nullptr);
    if (now > 1609459200) { // > Jan 1, 2021
      Serial.printf("[NTP] OK (unix=%ld)\n", (long)now);
      return;
    }
    delay(500);
  }
  Serial.println("[NTP] Time sync may have failed; continuing.");
}

bool probeDNS(const char* host, char outIp[16]) {
  struct addrinfo hints = {};
  hints.ai_family = AF_INET; // IPv4
  struct addrinfo* res = nullptr;

  Serial.printf("[DNS] Resolving %s...\n", host);
  int rc = getaddrinfo(host, NULL, &hints, &res);
  Serial.printf("[DNS] getaddrinfo rc=%d\n", rc);
  if (rc != 0 || !res) {
    Serial.println("[DNS] FAILED");
    return false;
  }
  struct sockaddr_in* sin = (struct sockaddr_in*)res->ai_addr;
  inet_ntop(AF_INET, &sin->sin_addr, outIp, 16);
  Serial.printf("[DNS] %s -> %s\n", host, outIp);
  freeaddrinfo(res);
  return true;
}


bool probeTCP(const char* host, uint16_t port, uint32_t timeout_ms = 5000) {
  WiFiClient cli;
  cli.setTimeout(timeout_ms);
  Serial.printf("[TCP] Connecting to %s:%u ...\n", host, port);
  if (!cli.connect(host, port)) {
    Serial.printf("[TCP] connect() FAILED\n");
    return false;
  }
  Serial.println("[TCP] Connected (no TLS). Closing (probe only).");
  cli.stop();
  return true;
}


bool probeTLS(const char* host, uint16_t port, uint32_t timeout_ms = 7000) {
  WiFiClientSecure tls;
  tls.setTimeout(timeout_ms);

  tls.setCACert(CA_ROOT_PEM);  

  Serial.printf("[TLS] Handshake to %s:%u ...\n", host, port);
  if (!tls.connect(host, port)) {
    int e = errno;
    Serial.printf("[TLS] connect() FAILED errno=%d (%s)\n", e, errnoName(e));
    return false;
  }
  Serial.println("[TLS] Handshake OK (server-only TLS)");
  tls.stop();
  return true;
}

bool probeMTLS(const char* host, uint16_t port, uint32_t timeout_ms = 8000) {
  WiFiClientSecure tls;
  tls.setTimeout(timeout_ms);

  tls.setCACert(CA_ROOT_PEM);
  tls.setCertificate(CLIENT_CERT_PEM);
  tls.setPrivateKey(CLIENT_KEY_PEM);

  Serial.printf("[mTLS] Handshake to %s:%u with client cert ...\n", host, port);
  if (!tls.connect(host, port)) {
    int e = errno;
    Serial.printf("[mTLS] connect() FAILED errno=%d (%s)\n", e, errnoName(e));
    Serial.println("[mTLS] If errno=ETIMEDOUT/ECONNRESET, server may be closing due to cert policy mismatch.");
    return false;
  }
  Serial.println("[mTLS] Handshake OK (client authenticated)");
  tls.stop();
  return true;
}

void setup() {
  Serial.begin(9600);
  delay(5000);
  Serial.println();
  Serial.println("==== Azure Event Grid MQTT Probe (mTLS) ====");

  WiFi.mode(WIFI_STA);

  if (!waitForWifi()) {
    Serial.println("[WiFi] FAILED to connect within timeout");
  } else {
    Serial.printf("[WiFi] Connected. IP=%s  RSSI=%d dBm\n",
                  WiFi.localIP().toString().c_str(), WiFi.RSSI());
  }

  // TLS sanity: time
  syncTime();

  // DNS
  char ip[16] = {0};
  bool dnsOk = probeDNS(HOSTNAME, ip);

  // TCP reachability
  bool tcpOk = probeTCP(HOSTNAME, PORT);

  // TLS (server-only)
  bool tlsOk = probeTLS(HOSTNAME, PORT);

  // TLS (mTLS with client cert/key)
  bool mtlsOk = probeMTLS(HOSTNAME, PORT);

  Serial.println("==== Summary ====");
  Serial.printf("DNS:  %s\n", dnsOk  ? "OK" : "FAILED");
  Serial.printf("TCP:  %s\n", tcpOk  ? "OK" : "FAILED");
  Serial.printf("TLS:  %s\n", tlsOk  ? "OK" : "FAILED");
  Serial.printf("mTLS: %s\n", mtlsOk ? "OK" : "FAILED");
  Serial.println("=================");

  Serial.println("If mTLS=FAILED, check: correct cert/key pair, chain/trust CA, and namespace mTLS policy.");
}

void loop() {
  delay(1000);
}