#include <WiFi.h>

const char* SSID = "";
const char* PASS = "";

uint8_t TARGET_BSSID[6];
int TARGET_CHANNEL = 0;

IPAddress STATIC_IP(192, 168, 1, 50);
IPAddress GATEWAY(192, 168, 1, 1);
IPAddress SUBNET(255, 255, 255, 0);
IPAddress DNS1(1, 1, 1, 1);
IPAddress DNS2(1, 0, 0, 1);

const int SAMPLES = 10;
const int NUM_CONFIGS = 16;
const IPAddress TARGET(1, 1, 1, 1);
const uint16_t PORT = 80;
const uint32_t CONNECT_TIMEOUT_MS = 5000;
const uint32_t WIFI_CONNECT_TIMEOUT_MS = 15000;

struct Result {
  char label[64];
  uint64_t wifiTotalMs;
  uint64_t tcpTotalMs;
  int count;
};

Result results[NUM_CONFIGS];

void halt(const char* msg) {
  Serial.print("FATAL: ");
  Serial.println(msg);
  while (true) delay(1000);
}

void discoverAP(const char* ssid) {
  Serial.println("Scanning for AP...");
  int n = WiFi.scanNetworks(false, true);
  if (n <= 0) halt("no networks found in scan");

  int bestIdx = -1, bestRSSI = -1000;
  for (int i = 0; i < n; i++) {
    if (strcmp(WiFi.SSID(i).c_str(), ssid) == 0 && WiFi.RSSI(i) > bestRSSI) {
      bestRSSI = WiFi.RSSI(i);
      bestIdx = i;
    }
  }
  if (bestIdx < 0) halt("target SSID not found in scan results");

  const uint8_t* bssidPtr = WiFi.BSSID(bestIdx);
  if (bssidPtr == nullptr) halt("BSSID pointer null - scan result corrupted");
  memcpy(TARGET_BSSID, bssidPtr, 6);
  TARGET_CHANNEL = WiFi.channel(bestIdx);
  WiFi.scanDelete();

  Serial.printf("Selected AP: BSSID=%s channel=%d RSSI=%d\n", WiFi.BSSIDstr(bestIdx).c_str(), TARGET_CHANNEL, bestRSSI);
}

// Full isolation between every single trial: driver stop/start clears
// PMK cache, BSSID/channel lock, and any lingering static-IP/DNS state.
// Costs more time per sample, but cost is identical for every trial,
// so relative comparisons between configs stay valid.
void resetRadio() {
  WiFi.disconnect(true, false); // drop connection, do NOT touch NVS
  WiFi.mode(WIFI_OFF);
  delay(100);
  WiFi.mode(WIFI_STA);
  delay(100);
}

void connectWiFi(bool useStatic, bool useChannel, bool useBSSID, bool useDNS, uint32_t &wifiMs, const char* label) {
  resetRadio();

  if (useStatic) {
    if (useDNS) WiFi.config(STATIC_IP, GATEWAY, SUBNET, DNS1, DNS2);
    else        WiFi.config(STATIC_IP, GATEWAY, SUBNET, GATEWAY); // explicit, not left to cache
  } else {
    if (useDNS) WiFi.config(INADDR_NONE, INADDR_NONE, INADDR_NONE, DNS1, DNS2); // DHCP IP, forced custom DNS
    else        WiFi.config((uint32_t)0, (uint32_t)0, (uint32_t)0);             // full DHCP, no override
  }

  uint32_t t0 = millis();
  int32_t channel = useChannel ? TARGET_CHANNEL : 0;
  uint8_t* bssid = useBSSID ? TARGET_BSSID : NULL;
  WiFi.begin(SSID, PASS, channel, bssid, true);

  while (WiFi.status() != WL_CONNECTED) {
    if (millis() - t0 > WIFI_CONNECT_TIMEOUT_MS) {
      char msg[96];
      snprintf(msg, sizeof(msg), "WiFi connect failed/timed out for config: %s", label);
      halt(msg);
    }
    delay(1);
  }
  wifiMs = millis() - t0;
}

uint32_t testTcpOnce(const char* label, int sampleNum) {
  WiFiClient client;
  uint32_t t0 = millis();
  bool connected = client.connect(TARGET, PORT, CONNECT_TIMEOUT_MS);
  uint32_t t1 = millis();
  if (!connected) {
    char msg[128];
    snprintf(msg, sizeof(msg), "TCP connect failed on sample %d for config: %s", sampleNum, label);
    halt(msg);
  }
  client.stop();
  return t1 - t0;
}

void buildLabel(int mask, char* out) {
  bool useStatic  = mask & 0x1, useChannel = mask & 0x2, useBSSID = mask & 0x4, useDNS = mask & 0x8;
  snprintf(out, 64, "IP:%s Chan:%s BSSID:%s DNS:%s",
           useStatic ? "static" : "dhcp", useChannel ? "fixed" : "auto",
           useBSSID ? "fixed" : "auto", useDNS ? "custom" : "default");
}

void setup() {
  Serial.begin(115200);
  delay(1000);
  randomSeed(esp_random());

  Serial.printf("=== ESP32-C3 single-AP config sweep (%d combos, %d samples each, isolated) ===\n", NUM_CONFIGS, SAMPLES);

  WiFi.persistent(false);
  WiFi.mode(WIFI_STA);
  WiFi.setSleep(false);
  discoverAP(SSID);

  for (int i = 0; i < NUM_CONFIGS; i++) buildLabel(i, results[i].label);

  // Flat list of (config, sample) pairs, shuffled once, so trial order
  // across the whole run is randomized rather than grouped/sequential.
  // Cancels out warm-radio / ordering bias across configs.
  int total = NUM_CONFIGS * SAMPLES;
  int order[total];
  for (int i = 0; i < total; i++) order[i] = i;
  for (int i = total - 1; i > 0; i--) {
    int j = random(i + 1);
    int tmp = order[i]; order[i] = order[j]; order[j] = tmp;
  }

  for (int k = 0; k < total; k++) {
    int mask = order[k] % NUM_CONFIGS;
    bool useStatic  = mask & 0x1, useChannel = mask & 0x2, useBSSID = mask & 0x4, useDNS = mask & 0x8;

    uint32_t wifiMs = 0;
    connectWiFi(useStatic, useChannel, useBSSID, useDNS, wifiMs, results[mask].label);
    uint32_t tcpMs = testTcpOnce(results[mask].label, results[mask].count + 1);

    results[mask].wifiTotalMs += wifiMs;
    results[mask].tcpTotalMs += tcpMs;
    results[mask].count++;

    Serial.printf("[%3d/%3d] %-64s wifi=%lu ms tcp=%lu ms\n", k + 1, total, results[mask].label, wifiMs, tcpMs);
  }

  Serial.println("\n=== Summary (sorted by avg WiFi connect time) ===");
  int idx[NUM_CONFIGS];
  for (int i = 0; i < NUM_CONFIGS; i++) idx[i] = i;
  for (int i = 0; i < NUM_CONFIGS - 1; i++)
    for (int j = i + 1; j < NUM_CONFIGS; j++)
      if (results[idx[i]].wifiTotalMs * results[idx[j]].count > results[idx[j]].wifiTotalMs * results[idx[i]].count) {
        int t = idx[i]; idx[i] = idx[j]; idx[j] = t;
      }

  Serial.println("Config                                                          | WiFi avg ms | TCP avg ms");
  for (int i = 0; i < NUM_CONFIGS; i++) {
    Result &r = results[idx[i]];
    Serial.printf("%-64s| %11.1f | %.1f\n", r.label, (float)r.wifiTotalMs / r.count, (float)r.tcpTotalMs / r.count);
  }
}

void loop() {}