#include <AccelStepper.h>
#include <Servo.h>

// === Stepper config for A4988 ===
#define STEP_PIN 8
#define DIR_PIN 9
AccelStepper myStepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

// === Servo setup ===
Servo servoRearArm;
Servo servoFrontArm;
const int servoRearPin = 4;
const int servoFrontPin = 3;

// === Joystick pins ===
const int X_pin_A = A0;
const int Y_pin_A = A1;
const int X_pin_B = A8;
const int Y_pin_B = A9;

// === Thresholds ===
const int THRESHOLD_LOW = 400;
const int THRESHOLD_HIGH = 600;
const int CENTER_THRESHOLD = 50;

// === Servo angles ===
int rearArmAngle = 90;
int frontArmAngle = 45;

// === Point memory ===
const int MAX_POINTS = 10;
String pointNames[MAX_POINTS];
int pointPositions[MAX_POINTS][3];
int pointCount = 0;

// === Trajectory recording ===
const int MAX_TRAJECTORY_POINTS = 50;
int trajectory[MAX_TRAJECTORY_POINTS][3];
int trajectoryLength = 0;
bool teaching = false;

// === Playback ===
bool movingToTarget = false;
bool playingBack = false;
bool playReverse = false;
int playbackIndex = 0;
long targetStepperPos = 0;
int targetRearArmAngle = 90;
int targetFrontArmAngle = 45;

// === For detecting pauses ===
unsigned long lastMoveTime = 0;
const unsigned long STOP_THRESHOLD = 600;

void setup() {
  Serial.begin(9600);
  Serial.println("serial poart");
  servoRearArm.attach(servoRearPin);
  servoFrontArm.attach(servoFrontPin);
  servoRearArm.write(rearArmAngle);
  servoFrontArm.write(frontArmAngle);
  myStepper.setMaxSpeed(1000.0);
  pinMode(X_pin_A, INPUT);
  pinMode(Y_pin_A, INPUT);
  pinMode(X_pin_B, INPUT);
  pinMode(Y_pin_B, INPUT);
}

void loop() {
  static int lastStepperPos = 0;
  static int lastRearAngle = 0;
  static int lastFrontAngle = 0;

  if (!movingToTarget) {
    checkJoystickA();
    checkJoystickB();
  }

  handleMoveToTarget();

  if (!movingToTarget && playingBack) {
    if (playbackIndex < 0 || playbackIndex >= trajectoryLength) {
      playingBack = false;
      Serial.println("Trajectory playback complete.");
    } else {
      targetStepperPos = trajectory[playbackIndex][0];
      targetRearArmAngle = trajectory[playbackIndex][1];
      targetFrontArmAngle = trajectory[playbackIndex][2];
      movingToTarget = true;
      if (playReverse)
        playbackIndex--;
      else
        playbackIndex++;
    }
  }

  myStepper.run();
  handleSerialCommands();

  // === Teach: detect stops to record trajectory points ===
  if (teaching) {
    bool isMoving = false;
    if (abs(myStepper.currentPosition() - lastStepperPos) > 3 ||
        abs(rearArmAngle - lastRearAngle) > 1 ||
        abs(frontArmAngle - lastFrontAngle) > 1) {
      isMoving = true;
      lastMoveTime = millis();
      lastStepperPos = myStepper.currentPosition();
      lastRearAngle = rearArmAngle;
      lastFrontAngle = frontArmAngle;
    }

    if (!isMoving && (millis() - lastMoveTime > STOP_THRESHOLD)) {
      if (trajectoryLength < MAX_TRAJECTORY_POINTS) {
        trajectory[trajectoryLength][0] = myStepper.currentPosition();
        trajectory[trajectoryLength][1] = rearArmAngle;
        trajectory[trajectoryLength][2] = frontArmAngle;
        trajectoryLength++;
        Serial.print("Recorded pause point #");
        Serial.println(trajectoryLength);
        lastMoveTime = millis() + 1000;
      }
    }
  }
}

void checkJoystickA() {
  int xValue = analogRead(X_pin_A);
  int centeredx = xValue - 512;
  int deadzone = 50;
  int yValue = analogRead(Y_pin_A);

  if (abs(centeredx) < deadzone) {
    myStepper.setSpeed(0);
  } else {
    float speed = (float)centeredx * 500.0 / 512.0;
    myStepper.setSpeed(speed);
    myStepper.runSpeed();
  }

  if (yValue > THRESHOLD_HIGH && rearArmAngle < 140) {
    rearArmAngle++;
    servoRearArm.write(rearArmAngle);
  } else if (yValue < THRESHOLD_LOW && rearArmAngle > 40) {
    rearArmAngle--;
    servoRearArm.write(rearArmAngle);
  }
}

void checkJoystickB() {
  int yValue = analogRead(Y_pin_B);
  if (yValue > THRESHOLD_HIGH && frontArmAngle < 180) {
    frontArmAngle++;
    servoFrontArm.write(frontArmAngle);
  } else if (yValue < THRESHOLD_LOW && frontArmAngle > 0) {
    frontArmAngle--;
    servoFrontArm.write(frontArmAngle);
  }
  delay(15);
}

void handleSerialCommands() {
  if (Serial.available()) {
    String input = Serial.readStringUntil('\n');
    input.trim();
    if (input.startsWith("remember(") && input.endsWith(")")) {
      String label = input.substring(9, input.length() - 1);
      remember(label);
    } else if (input.startsWith("moveto(") && input.endsWith(")")) {
      String label = input.substring(7, input.length() - 1);
      moveto(label);
    } else {
      Serial.println("Invalid command.");
    }
  }
}

void remember(String label) {
  if (pointCount < MAX_POINTS) {
    pointNames[pointCount] = label;
    pointPositions[pointCount][0] = myStepper.currentPosition();
    pointPositions[pointCount][1] = rearArmAngle;
    pointPositions[pointCount][2] = frontArmAngle;
    Serial.print("Recorded point '");
    Serial.print(label);
    Serial.println("' successfully.");

    if (!teaching) {
      trajectoryLength = 0;
      teaching = true;
      Serial.println("Teach mode STARTED.");
    } else {
      teaching = false;
      Serial.println("Teach mode ENDED. Path recorded.");
    }

    pointCount++;
  } else {
    Serial.println("Maximum number of points reached.");
  }
}

void moveto(String label) {
  for (int i = 0; i < pointCount; i++) {
    if (pointNames[i] == label) {
      int targetPos = pointPositions[i][0];
      int targetRear = pointPositions[i][1];
      int targetFront = pointPositions[i][2];

      if (!teaching && trajectoryLength > 0) {
        playReverse = (i == 0); // moveto(A) = reverse playback
        playbackIndex = playReverse ? trajectoryLength - 1 : 0;
        playingBack = true;
        Serial.println("Starting trajectory playback...");
        return;
      }

      // fallback: direct move
      targetStepperPos = targetPos;
      targetRearArmAngle = targetRear;
      targetFrontArmAngle = targetFront;
      movingToTarget = true;
      return;
    }
  }

  Serial.print("Point '");
  Serial.print(label);
  Serial.println("' not found.");
}

void handleMoveToTarget() {
  if (!movingToTarget) return;

  if (rearArmAngle < targetRearArmAngle) {
    rearArmAngle++;
    servoRearArm.write(rearArmAngle);
  } else if (rearArmAngle > targetRearArmAngle) {
    rearArmAngle--;
    servoRearArm.write(rearArmAngle);
  }

  if (frontArmAngle < targetFrontArmAngle) {
    frontArmAngle++;
    servoFrontArm.write(frontArmAngle);
  } else if (frontArmAngle > targetFrontArmAngle) {
    frontArmAngle--;
    servoFrontArm.write(frontArmAngle);
  }

  if (myStepper.distanceToGo() == 0 &&
      rearArmAngle == targetRearArmAngle &&
      frontArmAngle == targetFrontArmAngle) {
    movingToTarget = false;
    Serial.println("Reached target point.");
  }

  delay(10);
}