#include <WiFi.h>
#include <PubSubClient.h>

// Replace with your network credentials
const char* ssid = "New Boys Hostel";
const char* password = "nithostel%boys";

// Replace with your MQTT broker details
const char* mqtt_server = "test.mosquitto.org";
const char* control_topic = "train/braking/control";
const char* status_topic = "train/status/brakes";
const char* distance_topic = "train/distance";

// Initialize WiFi and MQTT clients
WiFiClient espClient;
PubSubClient client(espClient);

// Define HC-SR04 pins
const int trigPin = 5;
const int echoPin = 18;

// Define LED pin
const int ledPin = 2;

// Distance threshold in cm
const int distanceThreshold = 10;

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

  // Set up pins
  pinMode(trigPin, OUTPUT);
  pinMode(echoPin, INPUT);
  pinMode(ledPin, OUTPUT);

  // Connect to Wi-Fi
  setup_wifi();
  
  // Set up MQTT client
  client.setServer(mqtt_server, 1883);
  client.setCallback(callback);

  // Connect to MQTT broker
  connectMQTT();
}

void loop() {
  // Reconnect if disconnected
  if (!client.connected()) {
    connectMQTT();
  }
  
  client.loop();
  
  // Measure distance
  long duration, distance;
  digitalWrite(trigPin, LOW);
  delayMicroseconds(2);
  digitalWrite(trigPin, HIGH);
  delayMicroseconds(10);
  digitalWrite(trigPin, LOW);
  duration = pulseIn(echoPin, HIGH);
  distance = (duration / 2) * 0.0344; // Distance in cm

  Serial.print("Distance: ");
  Serial.print(distance);
  Serial.println(" cm");

  // Control LED based on distance
  if (distance < distanceThreshold) {
    digitalWrite(ledPin, HIGH);  // Turn LED on if too close
    client.publish(status_topic, "Object too close");
  } else {
    digitalWrite(ledPin, LOW);   // Turn LED off if safe
    client.publish(status_topic, "Object at safe distance");
  }

  // Publish distance data to MQTT topic
  String distanceString = String(distance);
  client.publish(distance_topic, distanceString.c_str());

  // Send periodic status updates
  static unsigned long lastUpdate = 0;
  if (millis() - lastUpdate > 10000) {  // Update every 10 seconds
    client.publish(status_topic, "Train operational");
    lastUpdate = millis();
  }
}

// Connect to Wi-Fi
void setup_wifi() {
  delay(10);
  Serial.println();
  Serial.print("Connecting to ");
  Serial.println(ssid);
  
  WiFi.begin(ssid, password);
  
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  
  Serial.println();
  Serial.println("Connected!");
  Serial.print("IP address: ");
  Serial.println(WiFi.localIP());
}

// Connect to MQTT broker
void connectMQTT() {
  while (!client.connected()) {
    Serial.print("Attempting MQTT connection...");
    
    // Attempt to connect
    if (client.connect("ESP32Client")) {
      Serial.println("connected");
      client.subscribe(control_topic);
    } else {
      Serial.print("failed, rc=");
      Serial.print(client.state());
      Serial.println(" try again in 5 seconds");
      delay(5000);
    }
  }
}

// Callback function for handling messages
void callback(char* topic, byte* payload, unsigned int length) {
  Serial.print("Message arrived [");
  Serial.print(topic);
  Serial.print("] ");
  
  String message;
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  
  Serial.println(message);
  
  if (message == "EMERGENCY") {
    Serial.println("Emergency brake activated!");
    client.publish(status_topic, "Brakes engaged");
  } else if (message == "STOP") {
    Serial.println("Train stopped.");
    client.publish(status_topic, "Train stopped");
  } else if (message == "RESUME") {
    Serial.println("Train resuming normal operation.");
    client.publish(status_topic, "Train resuming");
  }
}