#include <Arduino.h>
#include <ESP8266WiFi.h> // or #include <WiFi.h> for ESP32
#include <ESP8266mDNS.h> // or #include <ESPmDNS.h> for ESP32
#include <ESPAsyncWebServer.h>
#include <AsyncWebSocket.h>
#include <ArduinoJson.h>

// Node metadata structure
struct Node {
  String name;
  IPAddress ip;
  unsigned long lastSeen;
  String status;
};

// Service structure
struct Service {
  String name;
  String endpoint;
  void (*callback)();
};

// MemberList and ServiceRegistry
Node memberList[10]; // Example size
Service serviceRegistry[10];

AsyncWebServer server(80);
AsyncWebSocket ws("/ws");

void onServiceCall() {
  // Placeholder for service callback
}

// Helper to update node status
String getNodeStatus(unsigned long lastSeen) {
  unsigned long now = millis();
  if (now - lastSeen < 10000) return "active";
  if (now - lastSeen < 60000) return "inactive";
  return "dead";
}

// Discover other nodes via mDNS
void discoverNodes() {
  int n = MDNS.queryService("_ws", "tcp");
  if (n == 0) {
    Serial.println("No mDNS services found");
    return;
  }
  for (int i = 0; i < n && i < 10; i++) {
    IPAddress ip = MDNS.answerIP(i);
    String name = MDNS.answerHostname(i);
    if (ip == INADDR_NONE || name.isEmpty()) continue; // Skip invalid entries
    Serial.printf("Discovered node: %s (%s)\n", name.c_str(), ip.toString().c_str());
    // Add/update node in memberList
    bool found = false;
    for (int j = 0; j < 10; j++) {
      if (memberList[j].ip == ip) {
        memberList[j].lastSeen = millis();
        String oldStatus = memberList[j].status;
        memberList[j].status = getNodeStatus(memberList[j].lastSeen);
        found = true;
        if (oldStatus != memberList[j].status) {
          Serial.printf("Node %s (%s) status changed: %s -> %s\n", memberList[j].name.c_str(), memberList[j].ip.toString().c_str(), oldStatus.c_str(), memberList[j].status.c_str());
        }
        break;
      }
    }
    if (!found) {
      for (int j = 0; j < 10; j++) {
        if (memberList[j].name == "") {
          memberList[j].name = name;
          memberList[j].ip = ip;
          memberList[j].lastSeen = millis();
          memberList[j].status = "active";
          Serial.printf("Discovered new node: %s (%s)\n", name.c_str(), ip.toString().c_str());
          break;
        }
      }
    }
  }
}

// Broadcast event to all nodes
void broadcastEvent(const String& event) {
  DynamicJsonDocument doc(256);
  doc["type"] = "event";
  doc["data"] = event;
  String msg;
  serializeJson(doc, msg);
  ws.textAll(msg);
}

// Handle incoming WebSocket events
void onWsEvent(AsyncWebSocket *server, AsyncWebSocketClient *client, AwsEventType type, void *arg, uint8_t *data, size_t len) {
  if (type == WS_EVT_DATA) {
    String msg = String((char*)data);
    DynamicJsonDocument doc(256);
    DeserializationError err = deserializeJson(doc, msg);
    if (!err) {
      String type = doc["type"];
      if (type == "event") {
        // Handle pub/sub event
        Serial.println("Received event: " + doc["data"].as<String>());
      } else if (type == "service_call") {
        // Handle service call
        String endpoint = doc["endpoint"];
        for (int i = 0; i < 10; i++) {
          if (serviceRegistry[i].endpoint == endpoint && serviceRegistry[i].callback) {
            serviceRegistry[i].callback();
            break;
          }
        }
      }
    }
  }
}

extern "C" {
  #include <user_interface.h>
}

void onNewClient(System_Event_t *event) {
  if (event->event == EVENT_SOFTAPMODE_STACONNECTED) {
    uint8_t *mac = event->event_info.sta_connected.mac;
    Serial.printf("New WiFi client connected: MAC %02x:%02x:%02x:%02x:%02x:%02x\n",
      mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  }
}

String getMacHostname() {
  uint8_t mac[6];
  WiFi.macAddress(mac);
  char hostname[32];
  snprintf(hostname, sizeof(hostname), "esp-%02x%02x%02x%02x%02x%02x",
           mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  return String(hostname);
}

void setup() {
  Serial.begin(115200);
  delay(100);
  String myHostname = getMacHostname();
  WiFi.hostname(myHostname);
  WiFi.mode(WIFI_STA);
  WiFi.begin("contraption", "sprocket");
  Serial.printf("Attempting to connect to AP 'contraption' as %s...\n", myHostname.c_str());
  unsigned long startAttemptTime = millis();
  bool connected = false;
  while (millis() - startAttemptTime < 10000) { // try for 10 seconds
    if (WiFi.status() == WL_CONNECTED) {
      connected = true;
      break;
    }
    delay(500);
    Serial.print(".");
  }
  if (connected) {
    Serial.printf("\nConnected to AP 'contraption' as %s.\n", myHostname.c_str());
  } else {
    Serial.println("\nFailed to connect. Starting AP mode.");
    WiFi.mode(WIFI_AP);
    WiFi.softAP("contraption", "sprocket");
    WiFi.hostname(myHostname.c_str());
    Serial.printf("Access Point 'contraption' started with password 'sprocket' and hostname %s.\n", myHostname.c_str());
    wifi_set_event_handler_cb(onNewClient);
  }

  // mDNS setup
  if (MDNS.begin(myHostname)) {
    Serial.println("mDNS responder started");
    // Register WebSocket service on mDNS
    MDNS.addService("_ws", "tcp", 80); // 80 is the default WebSocket port
    Serial.println("mDNS WebSocket service '_ws._tcp' registered on port 80");
  }

  // WebSocket server setup
  ws.onEvent(onWsEvent);
  server.addHandler(&ws);
  server.begin();

  // Register example service
  serviceRegistry[0] = {"status", "/api/status", onServiceCall};
}

void loop() {
  // Periodically discover nodes
  static unsigned long lastDiscovery = 0;
  if (millis() - lastDiscovery > 5000) {
    discoverNodes();
    lastDiscovery = millis();
  }
  // Update node statuses
  for (int i = 0; i < 10; i++) {
    if (memberList[i].name != "") {
      String oldStatus = memberList[i].status;
      memberList[i].status = getNodeStatus(memberList[i].lastSeen);
      if (oldStatus != memberList[i].status) {
        Serial.printf("Node %s (%s) status changed: %s -> %s\n", memberList[i].name.c_str(), memberList[i].ip.toString().c_str(), oldStatus.c_str(), memberList[i].status.c_str());
      }
    }
  }
  // No need for webSocket.loop() with ESPAsyncWebServer
  // Node discovery, memberlist update, service registry logic to be added
}