#include "DHT.h"
#include "L298N.h"

#define DHTPIN 10
#define DHTTYPE DHT22
DHT dht(DHTPIN, DHTTYPE);

// --- Motor Setup ---
struct Motor {
  int stepPin;
  int dirPin;
  int speed;
  int direction;
  unsigned long lastStepTime;
};

Motor motors[5] = {
  {3, 2, 0, 1, 0},     // Motor 1
  {5, 4, 0, 1, 0},     // Motor 2
  {7, 6, 0, 1, 0},     // Motor 3 (Y-axis)
  {12, 11, 0, 1, 0},   // Motor 4 (Drum)
  {8, 9, 0, 1, 0}      // Motor 5 (External Stepper)
};

long motorPositions[5] = {0, 0, 0, 0, 0};  // Updated for 5 motors
bool motorHomed[5] = {false, false, false, false, false};  // Updated for 5 motors

const int limitSwitchPins[5] = {20, 21, 22, -1, 26}; // Updated for Motor 5 limit switch on pin 26

bool goToInitAfterHomeFlags[5] = {false, false, false, false, false};

// Loop movement settings
bool loopMovement = false;
int loopStart = 0;
int loopEnd = 150;
int yStepCount = 0;

// --- DHT Sensor ---
unsigned long lastDHTReadTime = 0;
const unsigned long dhtInterval = 5000;

// --- Relays ---
const int relayPins[4] = {13, 14, 15, 16};

// --- DC Motor ---
const int in1Pin = 17;
const int in2Pin = 18;
const int enaPin = 19;
int dcMotorSpeed = 255;

// --- INIT/END positions ---
const float mmPerStep = 5.0 / 200.0;
int initSteps = 0;
int endSteps = 150;

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

  for (int i = 0; i < 5; i++) {  // Update for 5 motors
    pinMode(motors[i].stepPin, OUTPUT);
    pinMode(motors[i].dirPin, OUTPUT);
    digitalWrite(motors[i].dirPin, motors[i].direction);
  }

  for (int i = 0; i < 4; i++) {
    pinMode(relayPins[i], OUTPUT);
    digitalWrite(relayPins[i], LOW);
  }

  pinMode(in1Pin, OUTPUT);
  pinMode(in2Pin, OUTPUT);
  pinMode(enaPin, OUTPUT);
  digitalWrite(in1Pin, LOW);
  digitalWrite(in2Pin, LOW);
  analogWrite(enaPin, 0);

  for (int i = 0; i < 5; i++) {  // Update for 5 motors
    if (limitSwitchPins[i] != -1) {
      pinMode(limitSwitchPins[i], INPUT_PULLUP);
    }
  }
}

void updateMotor(int motorId, int speed, int direction) {
  if (motorId < 1 || motorId > 5) return;  // Update to handle motor 5

  int idx = motorId - 1;
  
  // Ensure Motor is homed before any action
  if (motorHomed[idx] == false) {
    Serial.print("Motor ");
    Serial.print(motorId);
    Serial.println(" not homed. Please home it first.");
    return;
  }

  motors[idx].speed = speed;
  motors[idx].direction = direction;
  digitalWrite(motors[idx].dirPin, direction);

  // Pause loop if Motor 3 is manually controlled
  if (idx == 2 && loopMovement) {
    loopMovement = false;
    Serial.println("Loop paused due to manual control of Motor 3.");
  }
}

void updateDCMotor(int speed, int direction) {
  dcMotorSpeed = speed;
  digitalWrite(in1Pin, direction == 1 ? HIGH : LOW);
  digitalWrite(in2Pin, direction == 1 ? LOW : HIGH);
  analogWrite(enaPin, dcMotorSpeed);
}

void processCommand(String cmd) {
  int idx1 = cmd.indexOf(':');
  int idx2 = cmd.lastIndexOf(':');

  if (cmd.startsWith("M") && idx1 != -1 && idx2 != -1 && idx2 > idx1) {
    int motorId = cmd.substring(1, idx1).toInt();
    int speed = cmd.substring(idx1 + 1, idx2).toInt();
    int direction = cmd.substring(idx2 + 1).toInt();

    if (motorId == 3 && loopMovement) {
      loopMovement = false;
      Serial.println("Manual control detected. Looping disabled.");
    }

    updateMotor(motorId, speed, direction);
  }

  else if (cmd.startsWith("S") && idx1 != -1) {
    int motorId = cmd.substring(1, idx1).toInt();
    int speed = cmd.substring(idx1 + 1).toInt();
    int idx = motorId - 1;

    if (motorId >= 1 && motorId <= 5) {  // Updated for Motor 5
      if (!motorHomed[idx]) {
        Serial.print("Motor ");
        Serial.print(motorId);
        Serial.println(" not homed. Cannot set speed.");
        return;
      }

      motors[idx].speed = speed;
      Serial.print("Motor ");
      Serial.print(motorId);
      Serial.print(" speed updated to ");
      Serial.println(speed);
    }
  }

  else if (cmd.startsWith("LOOP:")) {
    int enable = cmd.substring(5).toInt();
    if (enable == 1) {
      if (!motorHomed[2]) {
        Serial.println("Motor 3 not homed. Cannot enable loop.");
        return;
      }
      if (motors[2].speed == 0) {
        motors[2].speed = 100;
        motors[2].direction = 1;
        digitalWrite(motors[2].dirPin, motors[2].direction);
        Serial.println("Default speed set for Motor 3 loop.");
      }
      loopMovement = true;
      Serial.println("Looping enabled.");
    } else {
        loopMovement = false;
        motors[2].speed = 0;  // Stop Motor 3
        Serial.println("Looping disabled. Motor 3 stopped.");
    }
  }

  else if (cmd.startsWith("INIT:")) {
    float mm = cmd.substring(5).toFloat();
    initSteps = mm / mmPerStep;
    Serial.print("INIT steps set to: ");
    Serial.println(initSteps);
  }

  else if (cmd.startsWith("END:")) {
    float mm = cmd.substring(4).toFloat();
    endSteps = mm / mmPerStep;
    Serial.print("END steps set to: ");
    Serial.println(endSteps);
  }

  else if (cmd.startsWith("HOME:")) {
    int motorId = cmd.substring(5).toInt();
    if (motorId >= 1 && motorId <= 5) {  // Updated for Motor 5
      int idx = motorId - 1;
      motors[idx].direction = 0;
      digitalWrite(motors[idx].dirPin, 0);
      motors[idx].speed = 100;
      Serial.print("Homing Motor ");
      Serial.println(motorId);
      
      while (digitalRead(limitSwitchPins[idx]) == HIGH) {  // Updated to use limit switch for Motor 5
        stepMotors();
      }
      motors[idx].speed = 0;
      motorPositions[idx] = 0;
      motorHomed[idx] = true;
      Serial.print("Motor ");
      Serial.print(motorId);
      Serial.println(" homed.");
    }
  }
}

void stepMotors() {
  unsigned long nowMicros = micros();

  for (int i = 0; i < 5; i++) {  // Update for 5 motors
    Motor& m = motors[i];
    if (m.speed > 0) {
      unsigned long interval = 1000000UL / m.speed;
      if (nowMicros - m.lastStepTime >= interval) {
        digitalWrite(m.stepPin, HIGH);
        delayMicroseconds(50);
        digitalWrite(m.stepPin, LOW);
        m.lastStepTime = nowMicros;

        m.direction == 1 ? motorPositions[i]++ : motorPositions[i]--;

        if (loopMovement && i == 2) {
          if (motorPositions[2] >= endSteps) {
            m.direction = 0;
            digitalWrite(m.dirPin, 0);
            Serial.println("Motor 3 reached END → reversing to INIT");
          }
          else if (motorPositions[2] <= initSteps) {
            m.direction = 1;
            digitalWrite(m.dirPin, 1);
            Serial.println("Motor 3 reached INIT → moving to END");
          }
        }
      }
    }
  }
}

void readDHTData() {
  float h = dht.readHumidity();
  float t = dht.readTemperature();
  float f = dht.readTemperature(true);

  if (isnan(h) || isnan(t) || isnan(f)) {
    Serial.println(F("Failed to read from DHT sensor!"));
    return;
  }

  Serial.print("Humidity: ");
  Serial.print(h);
  Serial.print("%  Temperature: ");
  Serial.print(t);
  Serial.print("°C ");
  Serial.println(f);
}

void loop() {
  stepMotors();

  if (millis() - lastDHTReadTime >= dhtInterval) {
    readDHTData();
    lastDHTReadTime = millis();
  }

  if (Serial.available()) {
    String input = Serial.readStringUntil('\n');
    Serial.println("Received: " + input);
    int start = 0;
    while (start < input.length()) {
      int end = input.indexOf(';', start);
      if (end == -1) end = input.length();
      String command = input.substring(start, end);
      command.trim();
      if (command.length() > 0) {
        processCommand(command);
      }
      start = end + 1;
    }
  }
}