#include <AccelStepper.h>
#include <ArduinoJson.h>

// ================= PIN =================
#define STEP_PIN    12
#define DIR_PIN     13
#define ENABLE_PIN  8
#define HOME_PIN    34

// ================= MOTORE =================
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

// ================= MECCANICA =================
const float passo_vite_mm = 12.0;
const int steps_per_rev = 200;
const int microstep = 16;
const float steps_per_mm = (steps_per_rev * microstep) / passo_vite_mm;

// ================= PARAMETRI =================
const float defaultMaxSpeed = 1200;   // step/s
const float accel = 600;              // step/s^2
const float homeSpeedFast = 600;      // step/s
const float homeBackoffSteps = 2 * steps_per_mm;

// ================= STATI =================
enum State {
  IDLE,
  HOME_FAST,
  HOME_BACKOFF,
  HOME_SLOW,
  MOVE_TO_POS,
  OSCILLATE
};

State state = IDLE;

// ================= OSCILLAZIONE =================
long centerPos = 0;
long deltaSteps = 0;
int totalCycles = 0;
int currentCycle = 0;
bool goingPositive = true;

// ================= FLAGS =================
bool homed = false;
bool stopRequested = false;

// ================= FEEDBACK =================
long lastPosReported = 0;

// =================================================
void setup() {
  Serial.begin(115200);

  pinMode(HOME_PIN, INPUT);
  pinMode(ENABLE_PIN, OUTPUT);
  digitalWrite(ENABLE_PIN, LOW);

  stepper.setAcceleration(accel);
  stepper.setMaxSpeed(defaultMaxSpeed);

  StaticJsonDocument<300> doc;
  doc["message"] = "Arduino UNO - Stepper FSM Ready";
  doc["commands"] = "HOME | MOVE pos_mm speed_mm_s | OSC delta_mm cicli speed_mm_s | STOP";
  serializeJson(doc, Serial);
  Serial.println();
}

// =================================================
void loop() {
  stepper.run();
  handleSerial();

  if (stopRequested) {
    emergencyStop();
    return;
  }

  switch (state) {
    case IDLE:           break;
    case HOME_FAST:      handleHomeFast(); break;
    case HOME_BACKOFF:   handleHomeBackoff(); break;
    case HOME_SLOW:      handleHomeSlow(); break;
    case MOVE_TO_POS:    handleMoveToPos(); break;
    case OSCILLATE:      handleOscillate(); break;
  }

  sendFeedback();
}

// ================= HANDLERS =================
void handleHomeFast() {
  if (digitalRead(HOME_PIN) == LOW) {
    stepper.move(homeBackoffSteps);
    state = HOME_BACKOFF;
  }
}

void handleHomeBackoff() {
  if (stepper.distanceToGo() == 0) {
    stepper.setMaxSpeed(homeSpeedFast / 4);
    stepper.moveTo(-10000);
    state = HOME_SLOW;
  }
}

void handleHomeSlow() {
  if (digitalRead(HOME_PIN) == LOW) {
    stepper.setCurrentPosition(0);
    stepper.setMaxSpeed(defaultMaxSpeed);
    homed = true;
    state = IDLE;

    StaticJsonDocument<200> doc;
    doc["event"] = "HOME_COMPLETED";
    serializeJson(doc, Serial);
    Serial.println();
  }
}

void handleMoveToPos() {
  if (stepper.distanceToGo() == 0) {
    stepper.setMaxSpeed(defaultMaxSpeed);
    state = IDLE;

    StaticJsonDocument<200> doc;
    doc["event"] = "MOVE_DONE";
    serializeJson(doc, Serial);
    Serial.println();
  }
}

void handleOscillate() {
  if (stepper.distanceToGo() == 0) {
    if (goingPositive) {
      stepper.moveTo(centerPos - deltaSteps);
      goingPositive = false;
    } else {
      currentCycle++;
      if (currentCycle >= totalCycles) {
        stepper.moveTo(centerPos);
        stepper.setMaxSpeed(defaultMaxSpeed);
        state = IDLE;

        StaticJsonDocument<200> doc;
        doc["event"] = "OSC_COMPLETED";
        serializeJson(doc, Serial);
        Serial.println();
      } else {
        stepper.moveTo(centerPos + deltaSteps);
        goingPositive = true;
      }
    }
  }
}

// ================= ERROR HELPER =================
void sendNotHomedError(const char* cmdName) {
  StaticJsonDocument<200> doc;
  doc["error"] = "NOT_HOMED";
  doc["command"] = cmdName;
  serializeJson(doc, Serial);
  Serial.println();
}

// ================= COMANDI SERIALI =================
void handleSerial() {
  if (!Serial.available()) return;

  String cmd = Serial.readStringUntil('\n');
  cmd.trim();

  if (cmd == "STOP") {
    stopRequested = true;
    return;
  }

  if (cmd == "HOME") {
    startHoming();
    return;
  }

  if (cmd.startsWith("MOVE")) {
    if (!homed) {
      sendNotHomedError("MOVE");
      return;
    }

    float pos, speed;
    if (sscanf(cmd.c_str(), "MOVE %f %f", &pos, &speed) == 2) {
      startMove(pos, speed);
    }
    return;
  }

  if (cmd.startsWith("OSC")) {
    if (!homed) {
      sendNotHomedError("OSC");
      return;
    }

    float delta, speed;
    int cicli;
    if (sscanf(cmd.c_str(), "OSC %f %d %f", &delta, &cicli, &speed) == 3) {
      startOscillation(delta, cicli, speed);
    }
    return;
  }
}

// ================= AZIONI =================
void startHoming() {
  if (state != IDLE) return;

  homed = false;
  stepper.setMaxSpeed(homeSpeedFast);
  stepper.moveTo(-100000);
  state = HOME_FAST;

  StaticJsonDocument<200> doc;
  doc["event"] = "HOME_STARTED";
  serializeJson(doc, Serial);
  Serial.println();
}

void startMove(float pos_mm, float speed_mm_s) {
  stepper.setMaxSpeed(speed_mm_s * steps_per_mm);
  centerPos = pos_mm * steps_per_mm;
  stepper.moveTo(centerPos);
  state = MOVE_TO_POS;
}

void startOscillation(float delta_mm, int cycles, float speed_mm_s) {
  stepper.setMaxSpeed(speed_mm_s * steps_per_mm);

  centerPos = stepper.currentPosition();
  deltaSteps = delta_mm * steps_per_mm;
  totalCycles = cycles;
  currentCycle = 0;
  goingPositive = true;

  stepper.moveTo(centerPos + deltaSteps);
  state = OSCILLATE;
}

// ================= EMERGENCY =================
void emergencyStop() {
  stepper.stop();
  stepper.runToPosition();
  digitalWrite(ENABLE_PIN, HIGH);

  stepper.setMaxSpeed(defaultMaxSpeed);
  state = IDLE;
  stopRequested = false;

  StaticJsonDocument<200> doc;
  doc["event"] = "EMERGENCY_STOP";
  serializeJson(doc, Serial);
  Serial.println();
}

// ================= FEEDBACK =================
void sendFeedback() {
  long pos = stepper.currentPosition();
  if (abs(pos - lastPosReported) >= steps_per_mm / 10) {
    StaticJsonDocument<200> doc;

    const char* stateName;
    switch (state) {
      case IDLE: stateName = "IDLE"; break;
      case HOME_FAST: stateName = "HOME_FAST"; break;
      case HOME_BACKOFF: stateName = "HOME_BACKOFF"; break;
      case HOME_SLOW: stateName = "HOME_SLOW"; break;
      case MOVE_TO_POS: stateName = "MOVE"; break;
      case OSCILLATE: stateName = "OSCILLATE"; break;
    }

    doc["state"] = stateName;
    doc["cycle"] = currentCycle;
    doc["pos_mm"] = pos / steps_per_mm;

    serializeJson(doc, Serial);
    Serial.println();
    lastPosReported = pos;
  }
}