#include <WiFi.h>
#include <DNSServer.h>
#include <WiFiClient.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

DNSServer dnsServer;
WiFiClient client;

const char* dnsProvidersIPv4[] = {
  "8.8.8.8",
  "94.140.14.14",
  "1.1.1.1",
  "9.9.9.9"
};

const char* dnsProvidersIPv6[] = {
  "2001:4860:4860::8888",
  "2a10:50c0::ad1:ff",
  "2606:4700:4700::1111",
  "2620:fe::fe"
};

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SSD1306_I2C_ADDRESS 0x3c
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire);

int totalQueries = 0;
int dnsLibraryCountIPv4[sizeof(dnsProvidersIPv4) / sizeof(dnsProvidersIPv4[0])];
int dnsLibraryCountIPv6[sizeof(dnsProvidersIPv6) / sizeof(dnsProvidersIPv6[0])];
int mostUsedLibraryIPv4 = -1;
int mostUsedLibraryIPv6 = -1;
int maxQueriesIPv4 = 0;
int maxQueriesIPv6 = 0;
int connectedClients = 0;

void setup() {
  Serial.begin(115200);
  initDisplay();
  connectToWiFi();
  dnsServer.start(53, "*", WiFi.localIP());
}

void loop() {
  dnsServer.processNextRequest();
  updateDisplay();
}

void initDisplay() {
  if (!display.begin(SSD1306_I2C_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
}

void connectToWiFi() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Connecting to WiFi...");

  WiFi.begin("Wokwi-GUEST", "");
  while (WiFi.status() != WL_CONNECTED) {
    delay(1000);
    Serial.print(".");
  }

  display.clearDisplay();
  display.setCursor(0, 0);
  if (WiFi.status() == WL_CONNECTED) {
    display.println("Connected to WiFi");
  } else {
    display.println("Failed to connect");
    for (;;);
  }
}

void resolveDNS(const char* domain, IPAddress& resolvedIP, const char* provider) {
  if (client.connect(provider, 53)) {
    // DNS query header
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x01);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x01);

    // Domain name parsing and query
    String parts[4];
    int i = 0;
    for (char c : String(provider)) {
      if (c == '.') {
        i++;
      } else {
        parts[i] += c;
      }
    }
    for (auto& part : parts) {
      client.write(part.toInt());
    }
    client.write(strlen(domain));
    client.write(domain);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x01);
    client.write((uint8_t)0x00);
    client.write((uint8_t)0x01);

    // Reduced delay
    delay(10);

    // Response handling
    while (client.available()) {
      // IPv6 address
      if (strstr(provider, ":")) {
        uint8_t responseLength = client.read();  // Read the length of the response
        char ipv6Buffer[40];
        client.readBytes(ipv6Buffer, responseLength);  // Read the response based on the length
        resolvedIP.fromString(ipv6Buffer);
      } else {
        resolvedIP = client.remoteIP();
      }
    }

    if (client.connected()) {
      if (strstr(provider, ":")) {
        updateDNSStatisticsIPv6(provider);
      } else {
        updateDNSStatisticsIPv4(provider);
      }
    }

    client.stop();
  }
}

void updateDNSStatisticsIPv4(const char* provider) {
  totalQueries++;
  updateMostUsedLibraryIPv4(provider);
  updateConnectedClients();
}

void updateDNSStatisticsIPv6(const char* provider) {
  totalQueries++;
  updateMostUsedLibraryIPv6(provider);
  updateConnectedClients();
}

void updateMostUsedLibraryIPv4(const char* provider) {
  for (int i = 0; i < sizeof(dnsProvidersIPv4) / sizeof(dnsProvidersIPv4[0]); i++) {
    if (strcmp(provider, dnsProvidersIPv4[i]) == 0) {
      if (++dnsLibraryCountIPv4[i] > maxQueriesIPv4) {
        maxQueriesIPv4 = dnsLibraryCountIPv4[i];
        mostUsedLibraryIPv4 = i;
      }
    }
  }
}

void updateMostUsedLibraryIPv6(const char* provider) {
  for (int i = 0; i < sizeof(dnsProvidersIPv6) / sizeof(dnsProvidersIPv6[0]); i++) {
    if (strcmp(provider, dnsProvidersIPv6[i]) == 0) {
      if (++dnsLibraryCountIPv6[i] > maxQueriesIPv6) {
        maxQueriesIPv6 = dnsLibraryCountIPv6[i];
        mostUsedLibraryIPv6 = i;
      }
    }
  }
}

void updateConnectedClients() {
  connectedClients = WiFi.softAPgetStationNum();
}

void updateDisplay() {
  display.clearDisplay();

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);

  display.print("Total Queries: ");
  display.println(totalQueries);

  display.print("Most Used Library (IPv4): ");
  if (mostUsedLibraryIPv4 != -1) {
    display.println(dnsProvidersIPv4[mostUsedLibraryIPv4]);
  } else {
    display.println("N/A");
  }

  display.print("Most Used Library (IPv6): ");
  if (mostUsedLibraryIPv6 != -1) {
    display.println(dnsProvidersIPv6[mostUsedLibraryIPv6]);
  } else {
    display.println("N/A");
  }

  display.print("Connected Clients: ");
  display.println(connectedClients);

  display.display();
}