#include <Arduino.h>
#include <TMCStepper.h>

// Define the number of motors
#define NUM_MOTORS 4

// Define default steps per command and step delay
// These are used to set how many steps the motors will attempt if started with a command.
const int defaultStepsPerCommand = 1600; // Adjust as needed for desired travel distance for a single command
// This is the period between steps for each motor in microseconds. Smaller value = faster movement.
const unsigned long stepPeriodMicroseconds = 1000; // 1000 us = 1ms -> 1000 steps/sec. Use higher values (e.g., 5000-10000) for slower, more visible testing.

// TMC2209 Driver configurations
#define R_SENSE 0.11f // Your driver's sense resistor value (e.g., 0.11 Ohm for most modules)

// TMC2209 motor driver parameters
const int runCurrent_mA = 800; // Motor run current in mA (max 1770mA with 0.11 ohm sense resistors)
const int holdCurrent_mA = 400; // Motor hold current in mA
const int microsteps = 16; // Number of microsteps (1, 2, 4, 8, 16, 32, 64, 128, 256)


// Define UART pins for the *single* shared UART for all motors.
// Based on the latest Wokwi JSON wiring, all TMC2209 drivers share UART2 on these specific pins.
#define M_ALL_UART_RX_PIN 16 // ESP32 GPIO 16 for shared UART RX
#define M_ALL_UART_TX_PIN 17 // ESP32 GPIO 17 for shared UART TX

// Define constants for specific motor pin assignments.
// These are matched to the 'motors' array initialization below and your Wokwi JSON wiring.
// IMPORTANT: Ensure these pin numbers match your Wokwi JSON exactly.
#define M1_STEP_PIN 18  // M1 (motors[0]) STEP pin
#define M1_DIR_PIN 19   // M1 (motors[0]) DIR pin
#define M1_DIAG_PIN 4   // M1 DIAG pin
#define M1_PDN_UART_PIN 5 // M1 PDN_UART/INDEX pin (often used as INDEX output or PDN_UART input)

#define M2_STEP_PIN 12  // M2 (motors[1]) STEP pin
#define M2_DIR_PIN 13   // M2 (motors[1]) DIR pin
#define M2_DIAG_PIN 34  // M2 DIAG pin
#define M2_PDN_UART_PIN 35 // M2 PDN_UART/INDEX pin

#define M3_STEP_PIN 14  // M3 (motors[2]) STEP pin
#define M3_DIR_PIN 27   // M3 (motors[2]) DIR pin
#define M3_DIAG_PIN 36  // M3 DIAG pin
#define M3_PDN_UART_PIN 39 // M3 PDN_UART/INDEX pin

#define M4_STEP_PIN 26  // M4 (motors[3]) STEP pin
#define M4_DIR_PIN 25   // M4 (motors[3]) DIR pin
#define M4_DIAG_PIN 21  // M4 DIAG pin
#define M4_PDN_UART_PIN 22 // M4 PDN_UART/INDEX pin


// Declare a single HardwareSerial instance for all motors.
// Using UART2 on ESP32, which defaults to GPIO 16/17.
HardwareSerial SerialMotors(2); // Using UART2 for ALL motors

// Motor structure to hold all relevant pins and driver object.
// The EN pin is removed from the struct as it's hardwired to GND in Wokwi.
struct Motor {
  uint8_t stepPin;
  uint8_t dirPin;
  uint8_t diagPin;
  uint8_t pdnUartPin;
  TMC2209Stepper *driver; // Pointer to the TMC2209Stepper object
  bool direction; // Current direction: HIGH or LOW (corresponds to digital state of DIR pin)
  bool moving; // Flag to indicate if the motor is currently commanded to move
  long stepsRemaining; // Number of steps left to take (for fixed-step movements)
  unsigned long lastStepTime; // Timestamp of the last step pulse (in microseconds)
  bool stepState; // Current state of the STEP pin (LOW or HIGH) for pulse generation
};

// Initialize motor array with GPIO pins as per your Wokwi JSON wiring.
Motor motors[NUM_MOTORS] = {
  // Motor 1 (array index 0)
  {M1_STEP_PIN, M1_DIR_PIN, M1_DIAG_PIN, M1_PDN_UART_PIN, nullptr, HIGH, false, 0, 0, LOW},
  // Motor 2 (array index 1) - Note: M2 often needs direction inverted for Mecanum
  {M2_STEP_PIN, M2_DIR_PIN, M2_DIAG_PIN, M2_PDN_UART_PIN, nullptr, LOW, false, 0, 0, LOW}, 
  // Motor 3 (array index 2)
  {M3_STEP_PIN, M3_DIR_PIN, M3_DIAG_PIN, M3_PDN_UART_PIN, nullptr, HIGH, false, 0, 0, LOW},
  // Motor 4 (array index 3)
  {M4_STEP_PIN, M4_DIR_PIN, M4_DIAG_PIN, M4_PDN_UART_PIN, nullptr, HIGH, false, 0, 0, LOW}
};


String receivedCommand = ""; // Buffer to store incoming serial data

// --- Function Declarations (Prototypes) ---
// These functions are defined later in the code.
void controlIndividualMotor(int motorIndex, bool forward);
void stopIndividualMotor(int motorIndex);
void controlAllMotors(String dirCommand);
void stopAllMotors();
void processCommand(String command);


// --- Functions Definitions ---

/**
 * @brief Controls a single motor's direction and starts it moving.
 * @param motorIndex The index of the motor in the 'motors' array (0-3).
 * @param forward Boolean: true for forward, false for backward.
 */
void controlIndividualMotor(int motorIndex, bool forward) {
  if (motorIndex < 0 || motorIndex >= NUM_MOTORS) {
    Serial.println("Invalid motor index.");
    return;
  }

  // Set initial direction (M2 is often inverted for Mecanum)
  motors[motorIndex].direction = (motorIndex == 1) ? !forward : forward;
  
  motors[motorIndex].moving = true; // Start motor movement
  // stepsRemaining is used for fixed-step movements, currently not actively stopping continuous movement
  motors[motorIndex].stepsRemaining = defaultStepsPerCommand; 
  motors[motorIndex].lastStepTime = micros(); // Record start time for step pulse generation
  motors[motorIndex].stepState = LOW; // Ensure step pin starts LOW for the first pulse

  // When a motor starts moving, ensure the driver is enabled and in SpreadCycle.
  // These settings are configured once in setup, but re-asserting can be useful.
  if (motors[motorIndex].driver) {
    motors[motorIndex].driver->toff(4); // Re-enable driver (if it was previously in low power mode)
    motors[motorIndex].driver->en_spreadCycle(true); // Force SpreadCycle (or StealthChop based on preference)
  }
  Serial.print("Motor "); Serial.print(motorIndex + 1); Serial.println(forward ? " starting forward" : " starting backward");
}

/**
 * @brief Stops a single motor.
 * @param motorIndex The index of the motor in the 'motors' array (0-3).
 */
void stopIndividualMotor(int motorIndex) {
  if (motorIndex >= 0 && motorIndex < NUM_MOTORS) {
    motors[motorIndex].moving = false; // Stop motor movement
    motors[motorIndex].stepsRemaining = 0; // Reset steps remaining
    digitalWrite(motors[motorIndex].stepPin, LOW); // Ensure step pin is low when stopped
    motors[motorIndex].stepState = LOW; // Reset step state
    Serial.print("Motor "); Serial.print(motorIndex + 1); Serial.println(" stopped");
    // Optionally: Put driver into low power hold mode
    // if (motors[motorIndex].driver) {
    //   motors[motorIndex].driver->toff(0); // Set toff to 0 to disable driver power stages
    // }
  }
}

/**
 * @brief Controls all motors for general chassis movement commands.
 * @param dirCommand String representing the desired movement ("forward", "backward", "left", "right", "turn left", "turn right").
 */
void controlAllMotors(String dirCommand) {
  // Always stop all motors first to reset state before new command
  stopAllMotors(); 

  // If the command was simply "stop", we're done after calling stopAllMotors()
  if (dirCommand == "stop") {
    return;
  }

  // Loop through all motors to prepare for movement
  for (int i = 0; i < NUM_MOTORS; i++) {
    motors[i].moving = true; // Set all motors to move
    motors[i].stepsRemaining = defaultStepsPerCommand; // This will ensure some movement happens before a 'stop' is needed
    motors[i].lastStepTime = micros(); // Initialize last step time for pulse generation
    motors[i].stepState = LOW; // Start with step pin low for the first pulse

    // Re-enable drivers (if they were in low power mode)
    if (motors[i].driver) {
      motors[i].driver->toff(4);
      motors[i].driver->en_spreadCycle(true); // Ensure SpreadCycle for reliability
    }
  }

  // Set specific directions for each motor based on the desired chassis movement
  // Adjust these directions if your Mecanum wheel configuration requires a different setup.
  if (dirCommand == "forward") {
    motors[0].direction = HIGH; // M1: Forward
    motors[1].direction = LOW;  // M2: Forward (inverted in code)
    motors[2].direction = HIGH; // M3: Forward
    motors[3].direction = HIGH; // M4: Forward
    Serial.println("Chassis moving forward.");
  } else if (dirCommand == "backward") {
    motors[0].direction = LOW;  // M1: Backward
    motors[1].direction = HIGH; // M2: Backward (inverted in code)
    motors[2].direction = LOW;  // M3: Backward
    motors[3].direction = LOW;  // M4: Backward
    Serial.println("Chassis moving backward.");
  } else if (dirCommand == "left") { // Strafe Left
    motors[0].direction = LOW;  // M1: Backward
    motors[1].direction = LOW;  // M2: Forward (inverted in code)
    motors[2].direction = HIGH; // M3: Forward
    motors[3].direction = LOW;  // M4: Backward
    Serial.println("Chassis strafing left.");
  } else if (dirCommand == "right") { // Strafe Right
    motors[0].direction = HIGH; // M1: Forward
    motors[1].direction = HIGH; // M2: Backward (inverted in code)
    motors[2].direction = LOW;  // M3: Backward
    motors[3].direction = HIGH; // M4: Forward
    Serial.println("Chassis strafing right.");
  } else if (dirCommand == "turn left") { // Rotate Counter-Clockwise
    motors[0].direction = LOW;  // M1: Backward
    motors[1].direction = LOW;  // M2: Forward (inverted in code)
    motors[2].direction = LOW;  // M3: Backward
    motors[3].direction = LOW;  // M4: Backward
    Serial.println("Chassis turning left.");
  } else if (dirCommand == "turn right") { // Rotate Clockwise
    motors[0].direction = HIGH; // M1: Forward
    motors[1].direction = HIGH; // M2: Backward (inverted in code)
    motors[2].direction = HIGH; // M3: Forward
    motors[3].direction = HIGH; // M4: Forward
    Serial.println("Chassis turning right.");
  } else {
    Serial.println("Unknown direction command for all motors: " + dirCommand);
    stopAllMotors(); // Stop if command is unrecognized
  }
}

/**
 * @brief Stops all motors in the chassis.
 */
void stopAllMotors() {
  for (int i = 0; i < NUM_MOTORS; i++) {
    motors[i].moving = false; // Stop motor movement
    motors[i].stepsRemaining = 0; // Reset steps
    digitalWrite(motors[i].stepPin, LOW); // Ensure step pin is low
    motors[i].stepState = LOW; // Reset step state
    // Optionally: Put driver into low power hold mode
    // if (motors[i].driver) {
    //   motors[i].driver->toff(0); // Set toff to 0 to disable driver power stages for low power hold
    // }
  }
  Serial.println("All motors stopped.");
}

/**
 * @brief Parses incoming serial commands and dispatches them to appropriate control functions.
 * @param command The string command received from the serial monitor.
 */
void processCommand(String command) {
  command.trim(); // Remove leading/trailing whitespace
  command.toLowerCase(); // Convert to lowercase for case-insensitive comparison

  Serial.print("Processing command: '");
  Serial.print(command);
  Serial.println("'");

  if (command == "stop") { // Universal stop command
    stopAllMotors();
  } else if (command.startsWith("m")) { // Individual motor control: "mX <direction>"
    int firstSpace = command.indexOf(' ');
    if (firstSpace != -1) {
      String motorIdStr = command.substring(0, firstSpace); // e.g., "m1"
      String dirAction = command.substring(firstSpace + 1); // e.g., "forward", "backward", "stop"

      int motorIndex = -1;
      if (motorIdStr == "m1") motorIndex = 0;
      else if (motorIdStr == "m2") motorIndex = 1;
      else if (motorIdStr == "m3") motorIndex = 2;
      else if (motorIdStr == "m4") motorIndex = 3;

      if (motorIndex != -1) {
        if (dirAction == "forward") {
          controlIndividualMotor(motorIndex, true);
        } else if (dirAction == "backward") {
          controlIndividualMotor(motorIndex, false);
        } else if (dirAction == "stop") {
          stopIndividualMotor(motorIndex);
        } else {
          Serial.println("Invalid individual motor action. Use 'forward', 'backward', or 'stop'.");
        }
      } else {
        Serial.println("Invalid motor identifier. Use 'm1', 'm2', 'm3', or 'm4'.");
      }
    } else {
      Serial.println("Invalid motor command format. Use 'm<number> <action>'.");
    }
  } else if (command.startsWith("all ")) { // All motors control: "all <direction>"
    String dirCommand = command.substring(4); // "all " is 4 characters
    controlAllMotors(dirCommand);
  } else {
    Serial.println("Unrecognized command.");
    Serial.println("Available commands: 'stop', 'm<num> forward/backward/stop', 'all forward/backward/left/right/turn left/turn right'");
  }
}

void setup() {
  Serial.begin(115200); // Initialize Serial Monitor for debugging
  Serial.println("Starting Mecanum Chassis Control...");

  // Initialize the single HardwareSerial instance for ALL motors (UART2)
  SerialMotors.begin(115200, SERIAL_8N1, M_ALL_UART_RX_PIN, M_ALL_UART_TX_PIN);
  Serial.print("UART for motors initialized on RX: "); Serial.print(M_ALL_UART_RX_PIN);
  Serial.print(", TX: "); Serial.println(M_ALL_UART_TX_PIN);

  // Configure motor structures and pins
  for (int i = 0; i < NUM_MOTORS; i++) {
    pinMode(motors[i].stepPin, OUTPUT);
    pinMode(motors[i].dirPin, OUTPUT);
    // EN pin is hardwired to GND in Wokwi, so no direct control from ESP32 is needed.
    // If you add a controllable EN pin in the future, uncomment and set pinMode/digitalWrite.
    // pinMode(motors[i].enPin, OUTPUT);
    // digitalWrite(motors[i].enPin, LOW); // Pull Enable LOW to activate driver (assuming active-low EN)
    digitalWrite(motors[i].stepPin, LOW); // Ensure step pin starts low

    // Initialize TMC2209Stepper objects with the *single* shared UART and unique addresses
    // The 'i' here maps directly to the address bits based on your MS1/MS2 Wokwi wiring:
    // i=0 -> Address 0b00 (M1)
    // i=1 -> Address 0b01 (M2)
    // i=2 -> Address 0b10 (M3)
    // i=3 -> Address 0b11 (M4)
    motors[i].driver = new TMC2209Stepper(&SerialMotors, R_SENSE, i); 

    Serial.print("Initializing TMC2209 driver for Motor "); Serial.print(i + 1); Serial.print(" (Address 0b"); Serial.print(i, BIN); Serial.println(")...");
    
    motors[i].driver->begin(); // Starts the UART communication and driver detection

    // Verify TMC2209 connection via UART
    if (!motors[i].driver->test_connection()) {
      Serial.print("ERROR: TMC2209 driver for Motor "); Serial.print(i + 1);
      Serial.println(" is NOT communicating via UART! Check wiring (UART, MS1/MS2), power, and resistor on RX line.");
      Serial.println("Halting program until fixed.");
      while(1); // Halt for debugging if connection fails
    } else {
      Serial.print("TMC2209 driver for Motor "); Serial.print(i + 1); Serial.println(" initialized and communicating.");
    }

    // Configure TMC2209 Driver Settings
    motors[i].driver->toff(4);              // Driver disable time (4 is a common value, 0 disables driver completely)
    motors[i].driver->rms_current(runCurrent_mA); // Set RMS running current in mA
    motors[i].driver->ihold(holdCurrent_mA);      // Set RMS holding current in mA
    motors[i].driver->irun(runCurrent_mA);        // Ensure irun matches rms_current for clarity
    motors[i].driver->microsteps(microsteps); // Set microsteps resolution
    motors[i].driver->en_spreadCycle(true); // Enable SpreadCycle (smoother, but might be slightly noisier than StealthChop)
    // motors[i].driver->en_stealthChop(true); // Uncomment for StealthChop (very quiet, but might lose torque at high speeds)
    motors[i].driver->pwm_autoscale(true);  // Automatically scale PWM for better performance

    // Print back current and microsteps to confirm settings were applied
    Serial.print("Motor "); Serial.print(i + 1); Serial.print(" configured: Run "); Serial.print(motors[i].driver->irun());
    Serial.print("mA, Hold "); Serial.print(motors[i].driver->ihold());
    Serial.print("mA, Microsteps "); Serial.println(motors[i].driver->microsteps());
    // Removed get_stealthChop() as it's not a standard method in the library.
  }

  Serial.println("\nMotor control initialized. Send commands via Serial Monitor.");
  Serial.println("--- COMMANDS ---");
  Serial.println("  'stop'              : Stops all motors.");
  Serial.println("  'm<num> forward'    : Starts individual motor <num> (1-4) forward.");
  Serial.println("  'm<num> backward'   : Starts individual motor <num> (1-4) backward.");
  Serial.println("  'm<num> stop'       : Stops individual motor <num> (1-4).");
  Serial.println("  'all forward'       : Moves chassis forward.");
  Serial.println("  'all backward'      : Moves chassis backward.");
  Serial.println("  'all left'          : Strafes chassis left.");
  Serial.println("  'all right'         : Strafes chassis right.");
  Serial.println("  'all turn left'     : Rotates chassis left.");
  Serial.println("  'all turn right'    : Rotates chassis right.");
  Serial.println("----------------");
}

void loop() {
  unsigned long currentMicros = micros();

  // Read incoming serial data (non-blocking) for commands
  while (Serial.available()) {
    char inChar = Serial.read();
    receivedCommand += inChar;
    if (inChar == '\n') { // Process command when newline character is received
      processCommand(receivedCommand);
      receivedCommand = ""; // Clear the command buffer
    }
  }

  // Handle motor stepping for all active motors
  for (int i = 0; i < NUM_MOTORS; i++) {
    if (motors[i].moving) {
      // Set the direction pin based on the motor's direction flag
      digitalWrite(motors[i].dirPin, motors[i].direction);

      // Generate a step pulse (HIGH for half period, LOW for half period)
      // This creates a continuous stepping motion as long as motors[i].moving is true.
      if (motors[i].stepState == LOW && (currentMicros - motors[i].lastStepTime >= stepPeriodMicroseconds / 2)) {
        digitalWrite(motors[i].stepPin, HIGH);
        motors[i].stepState = HIGH;
        motors[i].lastStepTime = currentMicros;
      } else if (motors[i].stepState == HIGH && (currentMicros - motors[i].lastStepTime >= stepPeriodMicroseconds / 2)) {
        digitalWrite(motors[i].stepPin, LOW);
        motors[i].stepState = LOW;
        motors[i].lastStepTime = currentMicros;
        // If you want fixed steps, you would decrement motors[i].stepsRemaining here
        // and set motors[i].moving = false when motors[i].stepsRemaining becomes 0.
      }
    } else {
      // Ensure step pin is LOW when the motor is not moving to prevent unintended steps
      digitalWrite(motors[i].stepPin, LOW);
    }
  }
}