#include <Servo.h>
#include <LiquidCrystal.h>

Servo servos[6];
int servoPins[] = {2, 3, 4, 5, 6, 7};
int btnPins[] = {22, 23, 24, 25, 26, 27};

const int dirSwitchPin = 28;  // S1
const int modeSwitchPin = 42; // S2
const int teachSwitchPin = 43; // S3

const int homeBtnPin = 36;
const int stopBtnPin = 37;
const int greyBtn1 = 41; // A2
const int greyBtn2 = 40; // A1
const int saveBtn = 38;  // crno dugme
const int confirmBtn = 39; // belo dugme

const int ledGreen = 44;
const int ledYellow = 45;
const int ledRed = 46;

LiquidCrystal lcd(30, 31, 32, 33, 34, 35);

int servoPositions[6] = {90, 90, 90, 90, 90, 90};

bool directionRight = true;
bool stopRequested = false;
bool teachingMode = false;
bool programSelected = false;
bool auto1Selected = false;
bool auto2Selected = false;

int stepCounter = 0;
int currentProgram = 0;
int programSteps[2][100][6];
int programStepCount[2] = {0, 0};

void setup() {
  for (int i = 0; i < 6; i++) {
    pinMode(btnPins[i], INPUT_PULLUP);
    servos[i].attach(servoPins[i]);
    servos[i].write(servoPositions[i]);
  }

  pinMode(dirSwitchPin, INPUT);
  pinMode(modeSwitchPin, INPUT);
  pinMode(teachSwitchPin, INPUT);
  pinMode(homeBtnPin, INPUT_PULLUP);
  pinMode(stopBtnPin, INPUT_PULLUP);
  pinMode(greyBtn1, INPUT_PULLUP);
  pinMode(greyBtn2, INPUT_PULLUP);
  pinMode(saveBtn, INPUT_PULLUP);
  pinMode(confirmBtn, INPUT_PULLUP);

  pinMode(ledGreen, OUTPUT);
  pinMode(ledYellow, OUTPUT);
  pinMode(ledRed, OUTPUT);

  lcd.begin(16, 2);
  updateDisplay();
}

void loop() {
  bool modeAuto = digitalRead(modeSwitchPin) == HIGH;
  bool teachActive = digitalRead(teachSwitchPin) == HIGH;
  directionRight = digitalRead(dirSwitchPin) == HIGH;

  digitalWrite(ledGreen, modeAuto && !teachActive);
  digitalWrite(ledYellow, !modeAuto && !teachActive);
  digitalWrite(ledRed, teachActive);

  stopRequested = digitalRead(stopBtnPin) == LOW;

  if (teachActive && !modeAuto && !programSelected) {
    if (digitalRead(greyBtn2) == LOW) {
      currentProgram = 0;
      programSelected = true;
      stepCounter = 0;
    }
    if (digitalRead(greyBtn1) == LOW) {
      currentProgram = 1;
      programSelected = true;
      stepCounter = 0;
    }
  }

  if (teachActive && !modeAuto) {
    for (int i = 0; i < 6; i++) {
      if (digitalRead(btnPins[i]) == LOW) {
        moveServo(i, directionRight ? 1 : -1);
      }
    }

    if (digitalRead(saveBtn) == LOW && programSelected) {
      if (stepCounter < 100) {
        for (int i = 0; i < 6; i++) {
          programSteps[currentProgram][stepCounter][i] = servoPositions[i];
        }
        stepCounter++;
      }
      delay(500);
    }

    if (digitalRead(confirmBtn) == LOW && programSelected) {
      programStepCount[currentProgram] = stepCounter;
      programSelected = false;
      delay(500);
    }
  }

  if (!teachActive && modeAuto && !stopRequested) {
    if (digitalRead(greyBtn2) == LOW) {
      runProgram(0);
    }
    if (digitalRead(greyBtn1) == LOW) {
      runProgram(1);
    }
  }

  if (!teachActive && !modeAuto) {
    for (int i = 0; i < 6; i++) {
      if (digitalRead(btnPins[i]) == LOW) {
        moveServo(i, directionRight ? 1 : -1);
      }
    }
  }

  if (digitalRead(homeBtnPin) == LOW && !stopRequested) {
    moveToHomeWhileHeld();
  }

  updateDisplay();
}

void moveServo(int index, int direction) {
  int newPos = constrain(servoPositions[index] + direction, 0, 180);
  if (newPos != servoPositions[index]) {
    servoPositions[index] = newPos;
    servos[index].write(newPos);
    delay(10);
  }
}

void moveToHomeWhileHeld() {
  int target = 15;
  while (digitalRead(homeBtnPin) == LOW) {
    for (int i = 4; i >= 2; i--) {
      if (servoPositions[i] != target) {
        int step = (servoPositions[i] > target) ? -1 : 1;
        servoPositions[i] += step;
        servos[i].write(servoPositions[i]);
        delay(10);
      }
    }

    bool moved = false;
    if (servoPositions[1] != target) {
      int step = (servoPositions[1] > target) ? -1 : 1;
      servoPositions[1] += step;
      servos[1].write(servoPositions[1]);
      moved = true;
    }
    if (servoPositions[0] != target) {
      int step = (servoPositions[0] > target) ? -1 : 1;
      servoPositions[0] += step;
      servos[0].write(servoPositions[0]);
      moved = true;
    }

    if (!moved && servoPositions[2] == target && servoPositions[3] == target && servoPositions[4] == target) {
      break;
    }
    delay(10);
  }
}

void runProgram(int progIndex) {
  for (int s = 0; s < programStepCount[progIndex]; s++) {
    for (int i = 0; i < 6; i++) {
      servoPositions[i] = programSteps[progIndex][s][i];
      servos[i].write(servoPositions[i]);
    }
    delay(300);
    if (digitalRead(stopBtnPin) == LOW) return;
  }
}

void updateDisplay() {
  static String lastLine1 = "", lastLine2 = "";
  bool modeAuto = digitalRead(modeSwitchPin) == HIGH;
  bool teachActive = digitalRead(teachSwitchPin) == HIGH;
  directionRight = digitalRead(dirSwitchPin) == HIGH;

  String line1 = "Dir: ";
  line1 += (teachActive && modeAuto) ? (directionRight ? "+" : "-") : (directionRight ? "R" : "L");

  if (teachActive && !modeAuto) {
    String countStr = (stepCounter < 10 ? "0" : "") + String(stepCounter);
    while (line1.length() < 16 - countStr.length()) line1 += " ";
    line1 += countStr;
  }

  if (teachActive && modeAuto) {
    while (line1.length() < 13) line1 += " ";
    line1 += "K ";
  }

  String line2 = "";
  if (teachActive) {
    if (programSelected) {
      line2 = "Auto" + String(currentProgram + 1);
    } else {
      line2 = "Ready";
    }
    while (line2.length() < 15) line2 += " ";
    line2 += "T";
  } else if (!modeAuto) {
    line2 = "Home";
    while (line1.length() < 16) line1 += " ";
    line1[15] = 'H';
  } else {
    line2 = "Auto";
    while (line1.length() < 16) line1 += " ";
    if (programSelected)
      line1[15] = (currentProgram == 0 ? 'A' : 'A') + String(currentProgram + 1)[0];
    else
      line1[15] = 'A';
  }

  if (line1 != lastLine1 || line2 != lastLine2) {
    lcd.clear();
    lcd.setCursor(0, 0); lcd.print(line1);
    lcd.setCursor(0, 1); lcd.print(line2);
    lastLine1 = line1;
    lastLine2 = line2;
  }
}