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

// Pins for components
const int ledRunningPin = 13;  // LED1 - Train Running
const int ledStoppedPin = 12;  // LED2 - Train Stopped
const int buttonPin = 14;      // Emergency stop button
const int potPin = 34;         // Potentiometer for speed control
const int stepPin = 26;        // A4988 STEP pin
const int dirPin = 27;         // A4988 DIR pin

// WiFi and MQTT setup
const char* ssid = "New Boys Hostel";
const char* password = "nithostel%boys";
const char* mqtt_server = "broker.hivemq.com";  // Use your MQTT broker

WiFiClient espClient;
PubSubClient client(espClient);

// Global variables
bool emergencyStop = false;
int trainSpeed = 0;
unsigned long lastStatusUpdate = 0;
const long statusUpdateInterval = 5000;  // Update interval for status (milliseconds)

// Function to connect to WiFi
void setup_wifi() {
  delay(10);
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
  }
}

// Callback function for incoming MQTT messages (server to device)
void callback(char* topic, byte* payload, unsigned int length) {
  String message = "";
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }

  // Handle control messages from server
  if (String(topic) == "train/control") {
    if (message == "start") {
      digitalWrite(ledRunningPin, HIGH);
      digitalWrite(ledStoppedPin, LOW);
      emergencyStop = false;
      client.publish("train/status", "Train started"); // Confirm start
    } else if (message == "stop") {
      digitalWrite(ledRunningPin, LOW);
      digitalWrite(ledStoppedPin, HIGH);
      emergencyStop = true;
      client.publish("train/status", "Train stopped"); // Confirm stop
    } else if (message == "emergency_stop") {
      digitalWrite(ledRunningPin, LOW);
      digitalWrite(ledStoppedPin, HIGH);
      emergencyStop = true;
      client.publish("train/status", "Emergency stop triggered"); // Confirm emergency stop
    }
  }
}

// MQTT reconnection
void reconnect() {
  while (!client.connected()) {
    if (client.connect("ESP32TrainClient")) {
      client.subscribe("train/control"); // Subscribe to control commands from the server
    } else {
      delay(5000);
    }
  }
}

void setup() {
  // Initialize components
  pinMode(ledRunningPin, OUTPUT);
  pinMode(ledStoppedPin, OUTPUT);
  pinMode(buttonPin, INPUT_PULLUP);   // Enable internal pull-up for pushbutton
  pinMode(stepPin, OUTPUT);
  pinMode(dirPin, OUTPUT);

  // Set up WiFi and MQTT
  setup_wifi();
  client.setServer(mqtt_server, 1883);
  client.setCallback(callback);

  // Initialize the train as stopped
  digitalWrite(ledRunningPin, LOW);
  digitalWrite(ledStoppedPin, HIGH);
  client.publish("train/status", "System initialized and stopped");
}

void loop() {
  // Reconnect to MQTT if necessary
  if (!client.connected()) {
    reconnect();
  }
  client.loop();

  // Read potentiometer for speed control
  int potValue = analogRead(potPin);
  trainSpeed = map(potValue, 0, 4095, 100, 1000);  // Map to speed range (100-1000 microseconds)

  // Check if emergency stop button is pressed
  if (digitalRead(buttonPin) == LOW) {
    emergencyStop = true;
    digitalWrite(ledRunningPin, LOW);
    digitalWrite(ledStoppedPin, HIGH);
    client.publish("train/status", "Emergency stop activated by button");
    delay(1000);  // Add a small delay to debounce
  }

  // If the train is not in emergency stop mode, run the motor
  if (!emergencyStop) {
    digitalWrite(dirPin, HIGH);  // Set direction of stepper motor
    for (int i = 0; i < trainSpeed; i++) {
      digitalWrite(stepPin, HIGH);
      delayMicroseconds(500);
      digitalWrite(stepPin, LOW);
      delayMicroseconds(500);
    }

    // Periodic status update to MQTT server
    if (millis() - lastStatusUpdate > statusUpdateInterval) {
      client.publish("train/status", "Train is running");
      lastStatusUpdate = millis();
    }
  } else {
    client.publish("train/status", "Train is stopped");
  }
}