#include <Arduino.h>
#include "FastAccelStepper.h"
#include <HardwareSerial.h>
#include <TMCStepper.h>
#include <SPI.h>

// Motor Pins
#define STEP_PIN 8
#define DIR_PIN 7
#define ENABLE_PIN 9
#define STALLGUARD 3
#define RXD2 17
#define TXD2 16

// Switch Pins
#define ON_OFF_SWITCH 10   // HIGH = OFF, LOW = ON
#define UP_DOWN_SWITCH 11  // HIGH = Up, LOW = Down
#define E_STOP_SWITCH 12   // LOW = Emergency Stop active

#define R_SENSE 0.11f
#define DRIVER_ADDRESS 0b00

const long set_velocity = 30000;
const int set_acceleration = 8000;
const int set_current = 500;
const int set_stall = 10;
const long set_tcools = 300;
const int motor_microsteps = 64;

volatile bool stalled_motor = false;
static bool lastWasEStop = false;  // Track if last command was e-stop

#define SERIAL_PORT_2 Serial2
FastAccelStepperEngine engine = FastAccelStepperEngine();
FastAccelStepper *stepper = NULL;
TMC2209Stepper driver(&SERIAL_PORT_2, R_SENSE, DRIVER_ADDRESS);

void stalled_position() {
  stalled_motor = true;
  Serial.println("STALL DETECTED!");
}

void setup() {
  Serial.begin(115200);
  SERIAL_PORT_2.begin(115200);

  pinMode(ON_OFF_SWITCH, INPUT_PULLUP);
  pinMode(UP_DOWN_SWITCH, INPUT_PULLUP);
  pinMode(E_STOP_SWITCH, INPUT_PULLUP);

  pinMode(ENABLE_PIN, OUTPUT);
  pinMode(STALLGUARD, INPUT);
  attachInterrupt(digitalPinToInterrupt(STALLGUARD), stalled_position, RISING);

  driver.begin();
  driver.toff(4);
  driver.blank_time(24);
  driver.I_scale_analog(false);
  driver.internal_Rsense(false);
  driver.mstep_reg_select(true);
  driver.microsteps(motor_microsteps);
  driver.TPWMTHRS(0);
  driver.semin(0);
  driver.shaft(true);
  driver.en_spreadCycle(false);
  driver.pdn_disable(true);
  driver.VACTUAL(0);
  driver.rms_current(set_current);
  driver.SGTHRS(set_stall);
  driver.TCOOLTHRS(set_tcools);

  engine.init();
  stepper = engine.stepperConnectToPin(STEP_PIN);
  stepper->setDirectionPin(DIR_PIN);
  stepper->setEnablePin(ENABLE_PIN);
  stepper->setAutoEnable(true);
  stepper->setSpeedInHz(set_velocity);
  stepper->setAcceleration(set_acceleration);

  // Move to top (UP) using runBackward()
  Serial.println("Initializing: Moving to Top...");
  stalled_motor = false;
  stepper->runBackward();
  while (!stalled_motor) {
    if (digitalRead(E_STOP_SWITCH) == LOW || digitalRead(ON_OFF_SWITCH) == HIGH) {
      stepper->forceStop();
      Serial.println("Startup interrupted by E-Stop or On/Off.");
      if (digitalRead(E_STOP_SWITCH) == LOW) lastWasEStop = true;  // Set flag during setup
      break;
    }
    Serial.print("SG_RESULT: "); Serial.println(driver.SG_RESULT());
    Serial.print("TSTEP: "); Serial.println(driver.TSTEP());
    delay(10);
  }
  stepper->forceStop();
  Serial.println("Reached Top at startup. Awaiting instructions...");
}

void loop() {
  static bool last_up_down_state = digitalRead(UP_DOWN_SWITCH);  // Initial state
  bool current_up_down_state = digitalRead(UP_DOWN_SWITCH);

  Serial.print("Loop running. Current state: "); Serial.print(current_up_down_state); 
  Serial.print(", Last state: "); Serial.println(last_up_down_state);

  // E-Stop check - only process if last command wasn't e-stop
  if (digitalRead(E_STOP_SWITCH) == LOW && !lastWasEStop) {
    stepper->forceStop();
    stepper->disableOutputs();
    Serial.println("Emergency Stop Activated!");
    while (digitalRead(E_STOP_SWITCH) == LOW) {
      delay(100);
      Serial.println("Waiting for E-Stop release...");
    }
    Serial.println("Emergency Stop Released. Awaiting instructions...");
    lastWasEStop = true;  // Mark that this was an e-stop
    last_up_down_state = digitalRead(UP_DOWN_SWITCH);  // Reset state
    return;
  } else if (digitalRead(E_STOP_SWITCH) == LOW && lastWasEStop) {
    Serial.println("Ignoring repeated E-Stop press...");
    while (digitalRead(E_STOP_SWITCH) == LOW) {
      delay(100);  // Wait for release without acting
    }
    Serial.println("Repeated E-Stop released. Awaiting instructions...");
    return;
  }

  // On/Off check
  if (digitalRead(ON_OFF_SWITCH) == HIGH) {
    stepper->forceStop();
    stepper->disableOutputs();
    Serial.println("System OFF");
    while (digitalRead(ON_OFF_SWITCH) == HIGH) {
      delay(100);
      Serial.println("Waiting for On/Off to enable...");
    }
    Serial.println("System ON. Awaiting instructions...");
    lastWasEStop = false;  // Clear e-stop flag
    last_up_down_state = digitalRead(UP_DOWN_SWITCH);  // Reset state
    return;
  }

  // Check for switch change and move
  if (current_up_down_state != last_up_down_state) {
    Serial.println("Switch changed to " + String(current_up_down_state ? "UP" : "DOWN") + ". Starting motion...");
    if (current_up_down_state == HIGH) {
      moveUp();
    } else {
      moveDown();
    }
    lastWasEStop = false;  // Clear e-stop flag after successful move
    last_up_down_state = current_up_down_state;  // Update state after move
  } else {
    Serial.println("No switch change. Awaiting instructions...");
  }
}

void moveUp() {
  Serial.println("Moving Up...");
  stalled_motor = false;
  stepper->runBackward();

  while (!stalled_motor) {
    if (digitalRead(E_STOP_SWITCH) == LOW || digitalRead(ON_OFF_SWITCH) == HIGH) return;

    bool new_state = digitalRead(UP_DOWN_SWITCH);
    if (new_state == LOW) {
      stepper->forceStop();
      Serial.println("Direction switched to Down.");
      return;
    }
    Serial.print("SG_RESULT: "); Serial.println(driver.SG_RESULT());
    Serial.print("TSTEP: "); Serial.println(driver.TSTEP());
    delay(10);
  }
  stepper->forceStop();
  Serial.println("Stalled (Top or Obstacle). Awaiting instructions...");
}

void moveDown() {
  Serial.println("Moving Down...");
  stalled_motor = false;
  stepper->runForward();

  while (!stalled_motor) {
    if (digitalRead(E_STOP_SWITCH) == LOW || digitalRead(ON_OFF_SWITCH) == HIGH) return;

    bool new_state = digitalRead(UP_DOWN_SWITCH);
    if (new_state == HIGH) {
      stepper->forceStop();
      Serial.println("Direction switched to Up.");
      return;
    }
    Serial.print("SG_RESULT: "); Serial.println(driver.SG_RESULT());
    Serial.print("TSTEP: "); Serial.println(driver.TSTEP());
    delay(10);
  }
  stepper->forceStop();
  Serial.println("Stalled (Bottom or Obstacle). Awaiting instructions...");
}