#include <painlessMesh.h>
#include <DHT.h>
#include <PubSubClient.h>
#include <WiFi.h>
#include <ArduinoJson.h>

// Mesh settings
#define MESH_PREFIX "IoTMesh"
#define MESH_PASSWORD "IoTMeshPassword"
#define MESH_PORT 5555

// Wi-Fi and MQTT settings (for bridge node)
#define WIFI_SSID "YourWiFiSSID"
#define WIFI_PASSWORD "YourWiFiPassword"
#define MQTT_BROKER "raspberry_pi_ip" // Replace with Raspberry Pi IP
#define MQTT_PORT 1883
#define MQTT_TOPIC_PUB "wot/nodes/data"
#define MQTT_TOPIC_SUB "wot/nodes/cmd"

// DHT11 settings
#define DHTPIN 4 // Pin connected to DHT11
#define DHTTYPE DHT11
DHT dht(DHTPIN, DHTTYPE);

// Node settings
#define NODE_ID "Node1" // Unique ID for this node
bool isBridge = false; // Set to true for the bridge node

// Global variables
painlessMesh mesh;
WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);
Scheduler userScheduler;

// Function prototypes
void sendSensorData();
void receivedCallback(uint32_t from, String &msg);
void mqttCallback(char *topic, byte *payload, unsigned int length);

void setup() {
  Serial.begin(115200);
  dht.begin();

  // Initialize mesh
  mesh.setDebugMsgTypes(ERROR | STARTUP);
  mesh.init(MESH_PREFIX, MESH_PASSWORD, &userScheduler, MESH_PORT);
  mesh.onReceive(&receivedCallback);

Task taskSendSensorData(TASK_SECOND * 10, TASK_FOREVER, &sendSensorData);

  // Check if this is the bridge node
  if (isBridge) {
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    while (WiFi.status() != WL_CONNECTED) {
      delay(500);
      Serial.print(".");
    }
    Serial.println("WiFi connected");

    mqttClient.setServer(MQTT_BROKER, MQTT_PORT);
    mqttClient.setCallback(mqttCallback);
    while (!mqttClient.connect("ESP32Bridge")) {
      Serial.print(".");
      delay(500);
    }
    Serial.println("MQTT connected");
    mqttClient.subscribe(MQTT_TOPIC_SUB);
  }

  // Schedule sensor data transmission
  userScheduler.addTask(taskSendSensorData);
  taskSendSensorData.enable();
}

void loop() {
  mesh.update();
  if (isBridge) {
    mqttClient.loop();
  }
}

void sendSensorData() {
  float temp = dht.readTemperature();
  float hum = dht.readHumidity();

  if (isnan(temp) || isnan(hum)) {
    Serial.println("Failed to read from DHT sensor!");
    return;
  }

  // Create JSON payload
  StaticJsonDocument<200> doc;
  doc["node"] = NODE_ID;
  doc["temperature"] = temp;
  doc["humidity"] = hum;
  String msg;
  serializeJson(doc, msg);

  // Send to mesh or MQTT
  if (isBridge) {
    mqttClient.publish(MQTT_TOPIC_PUB, msg.c_str());
  } else {
    mesh.sendBroadcast(msg);
  }

  Serial.printf("Sent: %s\n", msg.c_str());
}

void receivedCallback(uint32_t from, String &msg) {
  Serial.printf("Received from %u: %s\n", from, msg.c_str());

  // If bridge node, forward to MQTT
  if (isBridge) {
    mqttClient.publish(MQTT_TOPIC_PUB, msg.c_str());
  }
}

void mqttCallback(char *topic, byte *payload, unsigned int length) {
  String message;
  for (unsigned int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  Serial.printf("MQTT Message on %s: %s\n", topic, message.c_str());

  // Broadcast to mesh if needed
  mesh.sendBroadcast(message);
}