#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <AccelStepper.h>
#include <EEPROM.h>
#include <avr/wdt.h>

// Pin definitions
#define DIR_PIN 2
#define STEP_PIN 3
#define HOME_SENSOR_PIN 4
#define START_BUTTON_PIN 5
#define ESTOP_PIN 6
#define RELAY_PIN 7
#define UP_BUTTON_PIN 8
#define DOWN_BUTTON_PIN 9
#define RESET_PIN 10

LiquidCrystal_I2C lcd(0x27, 16, 2);
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

// EEPROM addresses
#define ADDR_DISTANCE 0
#define ADDR_PITCH 10
#define ADDR_MICROSTEP 20
#define ADDR_SPEED 30
#define ADDR_ACCEL 40
#define ADDR_RELAY_TIME 50
#define ADDR_PASSWORD 60

// System variables
int targetMM = 0;
float pitch = 8.0;
int microstep = 16;
float maxSpeed = 1000.0;
float acceleration = 500.0;
int relayTime = 2000;
float stepsPerMM;

int password[4];
int inputCode[4];
int codeIndex = 0;
bool passwordMode = false;
bool accessGranted = false;
bool inMenu = false;
int menuIndex = 0;

unsigned long resetHoldStart = 0;
bool resetTriggered = false;

void setup() {
  pinMode(HOME_SENSOR_PIN, INPUT_PULLUP);
  pinMode(START_BUTTON_PIN, INPUT_PULLUP);
  pinMode(ESTOP_PIN, INPUT_PULLUP);
  pinMode(RELAY_PIN, OUTPUT);
  pinMode(UP_BUTTON_PIN, INPUT_PULLUP);
  pinMode(DOWN_BUTTON_PIN, INPUT_PULLUP);
  pinMode(RESET_PIN, INPUT_PULLUP);

  lcd.begin(16, 2); // 16 columns, 2 rows
  lcd.backlight();

  // Watchdog setup
  wdt_enable(WDTO_2S);

  // Jumper-based password reset
  if (digitalRead(RESET_PIN) == LOW) {
    resetPasswordToDefault();
    lcd.print("Password Reset");
    delay(2000);
  }

  // Load settings
  EEPROM.get(ADDR_DISTANCE, targetMM);
  EEPROM.get(ADDR_PITCH, pitch);
  EEPROM.get(ADDR_MICROSTEP, microstep);
  EEPROM.get(ADDR_SPEED, maxSpeed);
  EEPROM.get(ADDR_ACCEL, acceleration);
  EEPROM.get(ADDR_RELAY_TIME, relayTime);
  loadPasswordFromEEPROM();

  calculateStepsPerMM();
  stepper.setMaxSpeed(maxSpeed);
  stepper.setAcceleration(acceleration);

  updateDisplay();
}

void loop() {
  wdt_reset(); // Reset watchdog

  checkLongPressReset();

  if (digitalRead(ESTOP_PIN) == LOW) {
    lcd.setCursor(0, 1);
    lcd.print("E-STOP Triggered ");
    stepper.stop();
    digitalWrite(RELAY_PIN, LOW);
    return;
  }

  if (passwordMode) {
    handlePasswordInput();
    return;
  }

  if (inMenu) {
    handleMenu();
    return;
  }

  if (digitalRead(START_BUTTON_PIN) == LOW && digitalRead(UP_BUTTON_PIN) == LOW) {
    enterPasswordMode();
    delay(500);
    return;
  }

  if (digitalRead(UP_BUTTON_PIN) == LOW) {
    targetMM++;
    updateDisplay();
    delay(200);
  }

  if (digitalRead(DOWN_BUTTON_PIN) == LOW) {
    if (targetMM > 0) targetMM--;
    updateDisplay();
    delay(200);
  }

  if (digitalRead(START_BUTTON_PIN) == LOW) {
    EEPROM.put(ADDR_DISTANCE, targetMM);
    lcd.setCursor(0, 1);
    lcd.print("Moving...        ");
    performHoming();
    moveToMM(targetMM);
    lcd.setCursor(0, 1);
    lcd.print("Done: ");
    lcd.print(targetMM);
    lcd.print("mm     ");
    digitalWrite(RELAY_PIN, HIGH);
    delay(relayTime);
    digitalWrite(RELAY_PIN, LOW);
  }
}

// ------------------ Password ------------------

void loadPasswordFromEEPROM() {
  for (int i = 0; i < 4; i++) {
    EEPROM.get(ADDR_PASSWORD + i, password[i]);
    if (password[i] < 0 || password[i] > 9) password[i] = i + 1;
  }
}

void resetPasswordToDefault() {
  int defaultPassword[4] = {1, 2, 3, 4};
  for (int i = 0; i < 4; i++) {
    EEPROM.put(ADDR_PASSWORD + i, defaultPassword[i]);
  }
}

void enterPasswordMode() {
  passwordMode = true;
  codeIndex = 0;
  for (int i = 0; i < 4; i++) inputCode[i] = 0;
  lcd.clear();
  lcd.print("Enter Password:");
  lcd.setCursor(0, 1);
  lcd.print("_ _ _ _");
}

void handlePasswordInput() {
  if (digitalRead(UP_BUTTON_PIN) == LOW) {
    inputCode[codeIndex]++;
    if (inputCode[codeIndex] > 9) inputCode[codeIndex] = 0;
    updatePasswordDisplay();
    delay(200);
  }

  if (digitalRead(DOWN_BUTTON_PIN) == LOW) {
    codeIndex++;
    if (codeIndex > 3) {
      checkPassword();
    }
    delay(200);
  }
}

void updatePasswordDisplay() {
  lcd.setCursor(0, 1);
  for (int i = 0; i < 4; i++) {
    if (i <= codeIndex) lcd.print(inputCode[i]);
    else lcd.print("_");
    lcd.print(" ");
  }
}

void checkPassword() {
  bool match = true;
  for (int i = 0; i < 4; i++) {
    if (inputCode[i] != password[i]) match = false;
  }

  if (match) {
    accessGranted = true;
    passwordMode = false;
    lcd.clear();
    lcd.print("Access Granted");
    delay(1000);
    inMenu = true;
    menuIndex = 0;
    showMenu();
  } else {
    lcd.clear();
    lcd.print("Wrong Password");
    delay(1000);
    passwordMode = false;
  }
}

void checkLongPressReset() {
  if (digitalRead(START_BUTTON_PIN) == LOW && digitalRead(DOWN_BUTTON_PIN) == LOW) {
    if (resetHoldStart == 0) resetHoldStart = millis();
    if (millis() - resetHoldStart > 5000 && !resetTriggered) {
      resetPasswordToDefault();
      lcd.clear();
      lcd.print("Password Reset");
      delay(2000);
      resetTriggered = true;
    }
  } else {
    resetHoldStart = 0;
    resetTriggered = false;
  }
}

// ------------------ Menu ------------------

void showMenu() {
  lcd.clear();
  switch (menuIndex) {
    case 0: lcd.print("Set Pitch:"); lcd.setCursor(0, 1); lcd.print(pitch, 1); break;
    case 1: lcd.print("Set uSteps:"); lcd.setCursor(0, 1); lcd.print(microstep); break;
    case 2: lcd.print("Set Speed:"); lcd.setCursor(0, 1); lcd.print(maxSpeed); break;
    case 3: lcd.print("Set Accel:"); lcd.setCursor(0, 1); lcd.print(acceleration); break;
    case 4: lcd.print("Relay Time:"); lcd.setCursor(0, 1); lcd.print(relayTime); break;
    case 5: lcd.print("Exit Config"); break;
  }
}

void handleMenu() {
  if (digitalRead(UP_BUTTON_PIN) == LOW) {
    adjustMenuValue(1);
    delay(200);
  }

  if (digitalRead(DOWN_BUTTON_PIN) == LOW) {
    adjustMenuValue(-1);
    delay(200);
  }

  if (digitalRead(START_BUTTON_PIN) == LOW) {
    confirmMenuSelection();
    delay(300);
  }
}

void adjustMenuValue(int delta) {
  switch (menuIndex) {
    case 0: pitch += delta * 0.1; if (pitch < 1.0) pitch = 1.0; if (pitch > 10.0) pitch = 10.0; break;
    case 1: microstep += delta; if (microstep < 1) microstep = 1; if (microstep > 32) microstep = 32; break;
    case 2: maxSpeed += delta * 50; if (maxSpeed < 100) maxSpeed = 100; if (maxSpeed > 3000) maxSpeed = 3000; break;
    case 3: acceleration += delta * 50; if (acceleration < 100) acceleration = 100; if (acceleration > 3000) acceleration = 3000; break;
    case 4: relayTime += delta * 500; if (relayTime < 500) relayTime = 500; if (relayTime > 10000) relayTime = 10000; break;
    case 5: menuIndex += delta; if (menuIndex < 0) menuIndex = 5; if (menuIndex > 5) menuIndex = 0; break;
  }
  showMenu();
}

void confirmMenuSelection() {
  if (menuIndex < 5) {
    menuIndex++;
    if (menuIndex > 5) menuIndex = 0;
    showMenu();
  } else {
    // Save settings
    EEPROM.put(ADDR_PITCH, pitch);
    EEPROM.put(ADDR_MICROSTEP, microstep);
    EEPROM.put(ADDR_SPEED, maxSpeed);
    EEPROM.put(ADDR_ACCEL, acceleration);
    EEPROM.put(ADDR_RELAY_TIME, relayTime);
    calculateStepsPerMM();
    stepper.setMaxSpeed(maxSpeed);
    stepper.setAcceleration(acceleration);
    inMenu = false;
    updateDisplay();
  }
}

// ------------------ Motion ------------------

void calculateStepsPerMM() {
  stepsPerMM = (200.0 * microstep) / pitch; // 200 = steps/rev for typical NEMA17/23
}

void performHoming() {
  stepper.setMaxSpeed(300);
  stepper.setAcceleration(200);
  while (digitalRead(HOME_SENSOR_PIN) == HIGH) {
    stepper.moveTo(stepper.currentPosition() - 1);
    stepper.run();
  }
  stepper.setCurrentPosition(0);
  stepper.setMaxSpeed(maxSpeed);
  stepper.setAcceleration(acceleration);
}

void moveToMM(int mm) {
  stepper.moveTo(mm * stepsPerMM);
  while (stepper.distanceToGo() != 0) {
    stepper.run();
  }
}

// ------------------ Display ------------------

void updateDisplay() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Target: ");
  lcd.print(targetMM);
  lcd.print(" mm");
}