#include <Arduino.h>
#include <math.h>

static const uint8_t GAUGE_COUNT = 4;

// For now, command and debug traffic share the same serial port.
#define CMD_PORT Serial
#define DEBUG_PORT Serial
static const unsigned long SERIAL_BAUD = 38400;

struct GaugePins {
  uint8_t dirPin;
  uint8_t stepPin;
  int8_t enablePin;          // -1 means there is no enable pin
  bool dirInverted;
  bool stepActiveHigh;
  bool enableActiveLow;
};

constexpr GaugePins gaugePins[GAUGE_COUNT] = {
  // dir, step, en, dirInv, stepHigh, enActiveLow
  {48, 49, -1, false, true, true},   // Gauge 0
  {8,  9,  -1, true,  true, true},   // Gauge 1
  {52, 53, -1, false, true, true},   // Gauge 2
  {50, 51, -1, false, true, true},   // Gauge 3
};

enum HomingState : uint8_t {
  HS_IDLE,
  HS_START,
  HS_BACKING,
  HS_SETTLE,
  HS_DONE
};

struct Gauge {
  long currentPos = 0;
  long targetPos = 0;

  long minPos = 0;
  long maxPos = 3780;
  long homingBackoffSteps = 3800; // Deliberately a touch past full reverse travel.

  float velocity = 0.0f;
  float maxSpeed = 4000.0f;
  float accel = 6000.0f;
  float homingSpeed = 500.0f;

  float stepAccumulator = 0.0f;
  unsigned long lastUpdateMicros = 0;

  bool enabled = true;
  bool homed = false;

  HomingState homingState = HS_IDLE;
  long homingStepsRemaining = 0;
  unsigned long homingLastStepMicros = 0;
  unsigned long homingStateStartMs = 0;

  bool sweepEnabled = false;
  bool sweepTowardMax = true;
};

Gauge gauges[GAUGE_COUNT];
String rxLine;

// Sends one-line command replies back over the control port.
//
// Serial protocol summary.
//
// Host -> controller commands (newline-terminated ASCII):
//   SET <id> <pos>
//   SPEED <id> <steps_per_s>
//   ACCEL <id> <steps_per_s2>
//   ENABLE <id> <0|1>
//   ZERO <id>
//   HOME <id>
//   HOMEALL
//   SWEEP <id> <accel> <speed>
//   POS?
//   CFG?
//   PING
//
// Controller -> host replies / events:
//   READY
//     Sent once from setup() after boot completes.
//   OK
//     Sent after a valid mutating command, and after POS?/CFG? once all data lines
//     for that query have been emitted.
//   PONG
//     Sent in response to PING.
//   ERR BAD_CMD
//     Sent when a complete line matches no parser.
//   ERR TOO_LONG
//     Sent when an input line exceeds the receive buffer limit.
//   ERR BAD_ID
//     Sent by commands that take a gauge id when the id is outside 0..GAUGE_COUNT-1.
//   ERR BAD_SPEED
//     Sent by SPEED when the requested speed is <= 0.
//   ERR BAD_ACCEL
//     Sent by ACCEL when the requested acceleration is <= 0.
//   POS <id> <currentPos> <targetPos> <homed> <homingState> <sweepEnabled>
//     Emitted once per gauge before the trailing OK reply to POS?.
//   CFG <id> <maxSpeed> <accel>
//     Emitted once per gauge before the trailing OK reply to CFG?.
//   HOMED <id>
//     Debug event printed on DEBUG_PORT when a homing sequence settles successfully.
void sendReply(const String& s) {
  CMD_PORT.println(s);
}

// Tiny float absolute-value helper to avoid dragging more machinery into the sketch.
float absf(float x) {
  return (x < 0.0f) ? -x : x;
}

// Updates the cached enable state and toggles the hardware pin if one exists.
void setEnable(uint8_t id, bool en) {
  if (id >= GAUGE_COUNT) return;
  gauges[id].enabled = en;

  int8_t pin = gaugePins[id].enablePin;
  if (pin < 0) return;

  bool level = gaugePins[id].enableActiveLow ? !en : en;
  digitalWrite(pin, level ? HIGH : LOW);
}

// Applies the logical direction after accounting for per-gauge inversion.
void setDir(uint8_t id, bool forward) {
  bool level = gaugePins[id].dirInverted ? !forward : forward;
  digitalWrite(gaugePins[id].dirPin, level ? HIGH : LOW);
}

// Emits one step pulse with the polarity expected by the driver.
void pulseStep(uint8_t id) {
  bool active = gaugePins[id].stepActiveHigh;
  digitalWrite(gaugePins[id].stepPin, active ? HIGH : LOW);
  delayMicroseconds(4);
  digitalWrite(gaugePins[id].stepPin, active ? LOW : HIGH);
}

// Moves the motor by one step if the requested direction is still within allowed travel.
void doStep(uint8_t id, int dir, bool allowPastMin = false) {
  Gauge& g = gauges[id];
  if (!g.enabled) return;

  if (dir > 0) {
    if (g.currentPos >= g.maxPos) return;
    setDir(id, true);
    pulseStep(id);
    g.currentPos++;
  } else if (dir < 0) {
    if (!allowPastMin && g.currentPos <= g.minPos) return;
    setDir(id, false);
    pulseStep(id);
    g.currentPos--;
  }
}

// Arms the homing state machine for one gauge and clears any in-flight motion.
void requestHome(uint8_t id) {
  if (id >= GAUGE_COUNT) return;
  Gauge& g = gauges[id];
  g.homingState = HS_START;
  g.homed = false;
  g.velocity = 0.0f;
  g.stepAccumulator = 0.0f;
  g.sweepEnabled = false;
}

// Starts the same homing sequence on every configured gauge.
void requestHomeAll() {
  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    requestHome(i);
  }
}

// Advances the simple homing state machine until the gauge is parked at logical zero.
void updateHoming(uint8_t id) {
  Gauge& g = gauges[id];
  unsigned long nowUs = micros();
  unsigned long nowMs = millis();

  switch (g.homingState) {
    case HS_IDLE:
      return;

    case HS_START:
      // No endstop here; homing just walks back far enough to hit the hard stop.
      g.velocity = 0.0f;
      g.stepAccumulator = 0.0f;
      g.homingStepsRemaining = g.homingBackoffSteps;
      g.homingLastStepMicros = nowUs;
      g.homingState = HS_BACKING;
      break;

    case HS_BACKING: {
      float intervalUs = 1000000.0f / g.homingSpeed;
      if ((nowUs - g.homingLastStepMicros) >= intervalUs) {
        g.homingLastStepMicros = nowUs;

        if (g.homingStepsRemaining > 0) {
          doStep(id, -1, true);
          g.homingStepsRemaining--;
        } else {
          g.homingState = HS_SETTLE;
          g.homingStateStartMs = nowMs;
        }
      }
      break;
    }

    case HS_SETTLE:
      if (nowMs - g.homingStateStartMs >= 100) {
        g.currentPos = 0;
        g.targetPos = 0;
        g.velocity = 0.0f;
        g.stepAccumulator = 0.0f;
        g.homed = true;
        g.homingState = HS_DONE;

        DEBUG_PORT.print("HOMED ");
        DEBUG_PORT.println(id);
      }
      break;

    case HS_DONE:
      g.homingState = HS_IDLE;
      break;
  }
}

// Flips the sweep destination when the gauge has settled at either end of travel.
void updateSweepTarget(uint8_t id) {
  Gauge& g = gauges[id];
  if (!g.sweepEnabled || !g.homed || g.homingState != HS_IDLE) return;

  if (g.sweepTowardMax) {
    g.targetPos = g.maxPos;
    if (g.currentPos >= g.maxPos && absf(g.velocity) < 1.0f) {
      g.sweepTowardMax = false;
      g.targetPos = g.minPos;
    }
  } else {
    g.targetPos = g.minPos;
    if (g.currentPos <= g.minPos && absf(g.velocity) < 1.0f) {
      g.sweepTowardMax = true;
      g.targetPos = g.maxPos;
    }
  }
}

// Runs one gauge worth of motion control, including homing and optional sweeping.
void updateGauge(uint8_t id) {
  Gauge& g = gauges[id];

  if (g.homingState != HS_IDLE) {
    updateHoming(id);
    return;
  }

  if (!g.homed) return;

  if (g.sweepEnabled) {
    updateSweepTarget(id);
  }

  unsigned long now = micros();
  if (g.lastUpdateMicros == 0) {
    g.lastUpdateMicros = now;
    return;
  }

  float dt = (now - g.lastUpdateMicros) / 1000000.0f;
  g.lastUpdateMicros = now;

  if (dt <= 0.0f || dt > 0.1f) return;

  long error = g.targetPos - g.currentPos;

  if (error == 0 && absf(g.velocity) < 0.01f) {
    g.velocity = 0.0f;
    g.stepAccumulator = 0.0f;
    return;
  }

  float dir = (error > 0) ? 1.0f : (error < 0 ? -1.0f : 0.0f);
  // Basic trapezoidal profile: brake if the remaining travel is shorter than the stop distance.
  float brakingDistance = (g.velocity * g.velocity) / (2.0f * g.accel + 0.0001f);

  if ((float)labs(error) <= brakingDistance) {
    if (g.velocity > 0.0f) {
      g.velocity -= g.accel * dt;
      if (g.velocity < 0.0f) g.velocity = 0.0f;
    } else if (g.velocity < 0.0f) {
      g.velocity += g.accel * dt;
      if (g.velocity > 0.0f) g.velocity = 0.0f;
    }
  } else {
    g.velocity += dir * g.accel * dt;
    if (g.velocity > g.maxSpeed) g.velocity = g.maxSpeed;
    if (g.velocity < -g.maxSpeed) g.velocity = -g.maxSpeed;
  }

  if (fabs(g.velocity) < 0.01f && error != 0) {
    g.velocity = dir * 5.0f;
  }

  // Integrate fractional steps until there is enough to emit a real pulse.
  g.stepAccumulator += g.velocity * dt;

  while (g.stepAccumulator >= 1.0f) {
    if (g.currentPos == g.targetPos) {
      g.stepAccumulator = 0.0f;
      g.velocity = 0.0f;
      break;
    }

    doStep(id, +1, false);
    g.stepAccumulator -= 1.0f;

    if (g.currentPos >= g.maxPos) {
      g.currentPos = g.maxPos;
      g.targetPos = g.maxPos;
      g.velocity = 0.0f;
      g.stepAccumulator = 0.0f;
      break;
    }
  }

  while (g.stepAccumulator <= -1.0f) {
    if (g.currentPos == g.targetPos) {
      g.stepAccumulator = 0.0f;
      g.velocity = 0.0f;
      break;
    }

    doStep(id, -1, false);
    g.stepAccumulator += 1.0f;

    if (g.currentPos <= g.minPos) {
      g.currentPos = g.minPos;
      g.targetPos = g.minPos;
      g.velocity = 0.0f;
      g.stepAccumulator = 0.0f;
      break;
    }
  }
}

// Parses `SET <id> <pos>` and updates the target position.
// Replies: `OK`, `ERR BAD_ID`.
bool parseSet(const String& line) {
  int id;
  long pos;
  if (sscanf(line.c_str(), "SET %d %ld", &id, &pos) == 2) {
    if (id < 0 || id >= GAUGE_COUNT) {
      sendReply("ERR BAD_ID");
      return true;
    }

    Gauge& g = gauges[id];
    pos = constrain(pos, g.minPos, g.maxPos);
    g.targetPos = pos;
    g.sweepEnabled = false;
    sendReply("OK");
    return true;
  }
  return false;
}

// Parses `SPEED <id> <speed>` and updates the max step rate.
// Replies: `OK`, `ERR BAD_ID`, `ERR BAD_SPEED`.
bool parseSpeed(const String& line) {
  int firstSpace = line.indexOf(' ');
  int secondSpace = line.indexOf(' ', firstSpace + 1);
  if (firstSpace < 0 || secondSpace < 0) return false;
  if (line.substring(0, firstSpace) != "SPEED") return false;

  int id = line.substring(firstSpace + 1, secondSpace).toInt();
  float speed = line.substring(secondSpace + 1).toFloat();

  if (id < 0 || id >= GAUGE_COUNT) {
    sendReply("ERR BAD_ID");
    return true;
  }
  if (speed <= 0.0f) {
    sendReply("ERR BAD_SPEED");
    return true;
  }

  gauges[id].maxSpeed = speed;
  sendReply("OK");
  return true;
}

// Parses `ACCEL <id> <accel>` and updates the acceleration limit.
// Replies: `OK`, `ERR BAD_ID`, `ERR BAD_ACCEL`.
bool parseAccel(const String& line) {
  int firstSpace = line.indexOf(' ');
  int secondSpace = line.indexOf(' ', firstSpace + 1);
  if (firstSpace < 0 || secondSpace < 0) return false;
  if (line.substring(0, firstSpace) != "ACCEL") return false;

  int id = line.substring(firstSpace + 1, secondSpace).toInt();
  float accel = line.substring(secondSpace + 1).toFloat();

  if (id < 0 || id >= GAUGE_COUNT) {
    sendReply("ERR BAD_ID");
    return true;
  }
  if (accel <= 0.0f) {
    sendReply("ERR BAD_ACCEL");
    return true;
  }

  gauges[id].accel = accel;
  sendReply("OK");
  return true;
}

// Parses `ENABLE <id> <0|1>` and toggles the selected driver.
// Replies: `OK`, `ERR BAD_ID`.
bool parseEnable(const String& line) {
  int id, en;
  if (sscanf(line.c_str(), "ENABLE %d %d", &id, &en) == 2) {
    if (id < 0 || id >= GAUGE_COUNT) {
      sendReply("ERR BAD_ID");
      return true;
    }

    setEnable(id, en != 0);
    sendReply("OK");
    return true;
  }
  return false;
}

// Parses `ZERO <id>` and declares the current position to be home.
// Replies: `OK`, `ERR BAD_ID`.
bool parseZero(const String& line) {
  int id;
  if (sscanf(line.c_str(), "ZERO %d", &id) == 1) {
    if (id < 0 || id >= GAUGE_COUNT) {
      sendReply("ERR BAD_ID");
      return true;
    }

    Gauge& g = gauges[id];
    g.currentPos = 0;
    g.targetPos = 0;
    g.velocity = 0.0f;
    g.stepAccumulator = 0.0f;
    g.homed = true;
    g.sweepEnabled = false;
    sendReply("OK");
    return true;
  }
  return false;
}

// Parses `HOME <id>` or `HOMEALL` and kicks off the homing sequence.
// Replies: `OK`, `ERR BAD_ID`. Successful completion later emits debug line `HOMED <id>`.
bool parseHome(const String& line) {
  int id;
  if (sscanf(line.c_str(), "HOME %d", &id) == 1) {
    if (id < 0 || id >= GAUGE_COUNT) {
      sendReply("ERR BAD_ID");
      return true;
    }

    requestHome(id);
    sendReply("OK");
    return true;
  }

  if (line == "HOMEALL") {
    requestHomeAll();
    sendReply("OK");
    return true;
  }

  return false;
}

// Parses `SWEEP <id> <accel> <speed>` and enables or disables end-to-end motion.
// Replies: `OK`, `ERR BAD_ID`.
bool parseSweep(const String& line) {
  int firstSpace = line.indexOf(' ');
  int secondSpace = line.indexOf(' ', firstSpace + 1);
  int thirdSpace = line.indexOf(' ', secondSpace + 1);

  if (firstSpace < 0 || secondSpace < 0 || thirdSpace < 0) return false;
  if (line.substring(0, firstSpace) != "SWEEP") return false;

  int id = line.substring(firstSpace + 1, secondSpace).toInt();
  float accel = line.substring(secondSpace + 1, thirdSpace).toFloat();
  float speed = line.substring(thirdSpace + 1).toFloat();

  if (id < 0 || id >= GAUGE_COUNT) {
    sendReply("ERR BAD_ID");
    return true;
  }

  Gauge& g = gauges[id];

  if (accel <= 0.0f || speed <= 0.0f) {
    g.sweepEnabled = false;
    g.velocity = 0.0f;
    g.stepAccumulator = 0.0f;
    sendReply("OK");
    return true;
  }

  g.accel = accel;
  g.maxSpeed = speed;
  g.sweepEnabled = true;
  g.sweepTowardMax = true;
  g.targetPos = g.maxPos;
  sendReply("OK");
  return true;
}

// Answers `POS?` with current motion state for every gauge.
// Emits one `POS <id> <cur> <tgt> <homed> <homingState> <sweep>` line per gauge,
// then replies `OK`.
bool parsePosQuery(const String& line) {
  if (line == "POS?") {
    for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
      CMD_PORT.print("POS ");
      CMD_PORT.print(i);
      CMD_PORT.print(' ');
      CMD_PORT.print(gauges[i].currentPos);
      CMD_PORT.print(' ');
      CMD_PORT.print(gauges[i].targetPos);
      CMD_PORT.print(' ');
      CMD_PORT.print(gauges[i].homed ? 1 : 0);
      CMD_PORT.print(' ');
      CMD_PORT.print((int)gauges[i].homingState);
      CMD_PORT.print(' ');
      CMD_PORT.println(gauges[i].sweepEnabled ? 1 : 0);
    }
    sendReply("OK");
    return true;
  }
  return false;
}

// Answers `CFG?` with speed and acceleration for every gauge.
// Emits one `CFG <id> <maxSpeed> <accel>` line per gauge, then replies `OK`.
bool parseCfgQuery(const String& line) {
  if (line == "CFG?") {
    for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
      CMD_PORT.print("CFG ");
      CMD_PORT.print(i);
      CMD_PORT.print(' ');
      CMD_PORT.print((int)gauges[i].maxSpeed);
      CMD_PORT.print(' ');
      CMD_PORT.println((int)gauges[i].accel);
    }
    sendReply("OK");
    return true;
  }
  return false;
}

// Answers the mandatory life question: are you there?
// Reply: `PONG`.
bool parsePing(const String& line) {
  if (line == "PING") {
    sendReply("PONG");
    return true;
  }
  return false;
}

// Runs the command parsers in order until one claims the line.
// Reply: `ERR BAD_CMD` when no parser accepts the line.
void processLine(const String& line) {
  if (parseSet(line)) return;
  if (parseSpeed(line)) return;
  if (parseAccel(line)) return;
  if (parseEnable(line)) return;
  if (parseZero(line)) return;
  if (parseHome(line)) return;
  if (parseSweep(line)) return;
  if (parsePosQuery(line)) return;
  if (parseCfgQuery(line)) return;
  if (parsePing(line)) return;

  sendReply("ERR BAD_CMD");
}

// Reads newline-delimited commands from serial and hands complete lines to the parser.
// Reply: `ERR TOO_LONG` when the buffered line exceeds the receive limit before newline.
void readCommands() {
  while (CMD_PORT.available()) {
    char c = (char)CMD_PORT.read();

    if (c == '\n') {
      rxLine.trim();
      if (rxLine.length() > 0) {
        processLine(rxLine);
      }
      rxLine = "";
    } else if (c != '\r') {
      rxLine += c;
      if (rxLine.length() > 120) {
        rxLine = "";
        sendReply("ERR TOO_LONG");
      }
    }
  }
}

// Initialises stepper pins and the initial homing cycle.
// Reply/event: emits `READY` on CMD_PORT once boot is complete.
void setup() {
  DEBUG_PORT.begin(SERIAL_BAUD);
  DEBUG_PORT.println("Gauge controller booting");

  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    pinMode(gaugePins[i].dirPin, OUTPUT);
    pinMode(gaugePins[i].stepPin, OUTPUT);

    digitalWrite(gaugePins[i].dirPin, LOW);
    digitalWrite(gaugePins[i].stepPin, gaugePins[i].stepActiveHigh ? LOW : HIGH);

    if (gaugePins[i].enablePin >= 0) {
      pinMode(gaugePins[i].enablePin, OUTPUT);
      setEnable(i, true);
    }

    gauges[i].lastUpdateMicros = micros();
  }

  requestHomeAll();

  DEBUG_PORT.println("READY");
  // Boot-complete handshake for the command channel.
  sendReply("READY");
}

// Main service loop: ingest commands and move gauges.
void loop() {
  readCommands();

  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    updateGauge(i);
  }
}