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

static const uint8_t GAUGE_COUNT = 2;

// For now: commands come over USB serial
#define CMD_PORT Serial
#define DEBUG_PORT Serial

struct GaugePins {
  uint8_t dirPin;
  uint8_t stepPin;
  int8_t enablePin;          // -1 if unused
  bool dirInverted;
  bool stepActiveHigh;
  bool enableActiveLow;
};

GaugePins gaugePins[GAUGE_COUNT] = {
  // dir, step, en, dirInv, stepHigh, enActiveLow
  {50, 51, -1, false, true, true},   // Gauge 0
  {8, 9, -1, true, true, true},   // Gauge 1
};

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 = 3610;             // adjust to your usable travel
  long homingBackoffSteps = 3700; // should exceed reverse travel slightly

  float velocity = 0.0f;
  float maxSpeed = 8000.0f;
  float accel = 9000.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;

void sendReply(const String& s) {
  CMD_PORT.println(s);
}

float absf(float x) {
  return (x < 0.0f) ? -x : x;
}

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);
}

void setDir(uint8_t id, bool forward) {
  bool level = gaugePins[id].dirInverted ? !forward : forward;
  digitalWrite(gaugePins[id].dirPin, level ? HIGH : LOW);
}

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);
}

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--;
  }
}

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;
}

void requestHomeAll() {
  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    requestHome(i);
  }
}

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:
      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;
  }
}

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;
    }
  }
}

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);
  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;
  }

  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;
    }
  }
}

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;
}

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;
}

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;
}

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;
}

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;
}

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;
}

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;
}

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);
    }
    return true;
  }
  return false;
}

bool parsePing(const String& line) {
  if (line == "PING") {
    sendReply("PONG");
    return true;
  }
  return false;
}

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 (parsePing(line)) return;

  sendReply("ERR BAD_CMD");
}

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 = "";
      }
    }
  }
}

void setup() {
  DEBUG_PORT.begin(115200);
  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");
  sendReply("READY");
}

void loop() {
  readCommands();

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