#include <WiFi.h>
#include <AccelStepper.h>
#include <WebServer.h>

// Define stepper motors and pins
AccelStepper stepper1(AccelStepper::DRIVER, 16, 15); // STEP pin 16, DIR pin 15
AccelStepper stepper2(AccelStepper::DRIVER, 14, 13); // STEP pin 14, DIR pin 13

// Joystick pins
const int VERT = 12;
const int HOR = 11;
const int PUSH = 10;

// Variables to store neutral (center) position of the joystick
int neutralX = 0;
int neutralY = 0;
bool calibrated = false;

// WiFi credentials
const char* ssid = "YOUR_SSID";
const char* password = "YOUR_PASSWORD";

// Create a web server object that listens for HTTP requests on port 80
WebServer server(80);

void setup() {
  // Initialize serial communication
  Serial.begin(115200);

  // Set up joystick button pin
  pinMode(PUSH, INPUT_PULLUP);

  // Set max speed and acceleration for stepper motors
  stepper1.setMaxSpeed(1000);
  stepper1.setAcceleration(500);
  stepper2.setMaxSpeed(1000);
  stepper2.setAcceleration(500);

  // Initial calibration
  calibrateJoystick();

  // Connect to WiFi
  connectToWiFi();

  // Define the endpoints for the API
  server.on("/control", handleMotorControl);
  
  // Start the server
  server.begin();
}

void loop() {
  // Handle any incoming client requests
  server.handleClient();

  // Joystick control routine
  joystickControl();
}

void connectToWiFi() {
  Serial.println("Connecting to WiFi...");
  WiFi.begin(ssid, password);

  while (WiFi.status() != WL_CONNECTED) {
    delay(1000);
    Serial.println("Connecting...");
  }

  Serial.println("Connected to WiFi!");
  Serial.print("IP Address: ");
  Serial.println(WiFi.localIP());
}

void calibrateJoystick() {
  Serial.println("Calibrating joystick...");

  // Wait for the button to be held down for 2 seconds to start calibration
  unsigned long startTime = millis();
  while (digitalRead(PUSH) == LOW) {
    if (millis() - startTime > 2000) {
      break;
    }
  }

  // Read the current joystick values as neutral positions
  neutralX = analogRead(HOR);
  neutralY = analogRead(VERT);

  Serial.print("Neutral X: ");
  Serial.println(neutralX);
  Serial.print("Neutral Y: ");
  Serial.println(neutralY);

  calibrated = true;
}

void joystickControl() {
  // Read joystick values
  int xVal = analogRead(HOR);
  int yVal = analogRead(VERT);

  // Adjust values based on calibration
  int adjustedX = xVal - neutralX;
  int adjustedY = yVal - neutralY;

  // Map adjusted joystick values to stepper speeds
  long speed1 = map(adjustedX, -2048, 2048, -1000, 1000);
  long speed2 = map(adjustedY, -2048, 2048, -1000, 1000);

  // Set speed for each motor only if there is a significant deviation from neutral
  if (abs(adjustedX) > 50) { // Deadzone threshold
    stepper1.setSpeed(speed1);
  } else {
    stepper1.setSpeed(0);
  }

  if (abs(adjustedY) > 50) { // Deadzone threshold
    stepper2.setSpeed(speed2);
  } else {
    stepper2.setSpeed(0);
  }

  // Move the motors
  stepper1.runSpeed();
  stepper2.runSpeed();
}

void handleMotorControl() {
  if (server.hasArg("Motor") && server.hasArg("Direction") && server.hasArg("Speed") && server.hasArg("Steps") && server.hasArg("Mode")) {
    int motor = server.arg("Motor").toInt();
    int direction = server.arg("Direction").toInt();
    int speed = server.arg("Speed").toInt();
    int steps = server.arg("Steps").toInt();
    String mode = server.arg("Mode");

    AccelStepper *stepper;
    if (motor == 1) {
      stepper = &stepper1;
    } else if (motor == 2) {
      stepper = &stepper2;
    } else {
      server.send(400, "text/plain", "Invalid motor selection");
      return;
    }

    if (direction == 1) {
      stepper->move(steps);
    } else if (direction == 0) {
      stepper->move(-steps);
    } else {
      server.send(400, "text/plain", "Invalid direction");
      return;
    }

    if (mode == "RAW") {
      stepper->setMaxSpeed(speed);
      stepper->setSpeed(speed);
      while (stepper->distanceToGo() != 0) {
        stepper->runSpeedToPosition();
      }
    } else if (mode == "SMOOTH") {
      smoothMove(stepper, speed, steps);
    } else {
      server.send(400, "text/plain", "Invalid mode");
      return;
    }

    server.send(200, "text/plain", "Motor control command executed");
  } else {
    server.send(400, "text/plain", "Missing required parameters");
  }
}

void smoothMove(AccelStepper *stepper, int speed, int steps) {
  stepper->setAcceleration(100); // Set acceleration for smooth start/stop
  stepper->setMaxSpeed(speed);
  stepper->move(steps);
  while (stepper->distanceToGo() != 0) {
    stepper->run();
  }
  stepper->setAcceleration(500); // Reset to original acceleration value
}