#include <Arduino.h>
#include <avr/interrupt.h>
#include <math.h>
#include "gauge_config.h"

static const uint16_t STEPPER_TIMER_HZ = 20000;
static const int32_t STEP_RATE_SCALE = 65536L;
static const float MAX_TIMER_STEP_RATE = (float)STEPPER_TIMER_HZ / 2.0f;

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

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

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

  long minPos = 0;
  long maxPos = 0;
  long homingBackoffSteps = 0;

  float velocity = 0.0f;
  float maxSpeed = 0.0f;
  float accel = 0.0f;
  float homingSpeed = 0.0f;

  unsigned long lastUpdateMicros = 0;

  volatile int32_t timerStepRateQ16 = 0;
  volatile int32_t timerStepAccumulatorQ16 = 0;
  volatile bool timerAllowPastMin = false;
  volatile bool timerPulseActive = false;

  volatile bool enabled = true;
  bool homed = false;

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

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

Gauge gauges[GAUGE_COUNT];
String rxLine;

struct StepperHardware {
  volatile uint8_t* stepPort = nullptr;
  volatile uint8_t* dirPort = nullptr;
  uint8_t stepMask = 0;
  uint8_t dirMask = 0;
};

StepperHardware stepperHardware[GAUGE_COUNT];

long atomicReadLong(volatile long& value) {
  uint8_t sreg = SREG;
  noInterrupts();
  long copy = value;
  SREG = sreg;
  return copy;
}

void atomicWriteLong(volatile long& value, long newValue) {
  uint8_t sreg = SREG;
  noInterrupts();
  value = newValue;
  SREG = sreg;
}

inline void writePortBit(volatile uint8_t* port, uint8_t mask, bool high) {
  if (high) {
    *port |= mask;
  } else {
    *port &= ~mask;
  }
}

inline void writeDirectionPin(uint8_t id, bool forward) {
  bool level = gaugeConfigs[id].dirInverted ? !forward : forward;
  writePortBit(stepperHardware[id].dirPort, stepperHardware[id].dirMask, level);
}

inline void writeStepPin(uint8_t id, bool active) {
  bool level = gaugeConfigs[id].stepActiveHigh ? active : !active;
  writePortBit(stepperHardware[id].stepPort, stepperHardware[id].stepMask, level);
}

inline void stepDirectionSetupDelay() {
  __asm__ __volatile__(
      "nop\n\t""nop\n\t""nop\n\t""nop\n\t"
      "nop\n\t""nop\n\t""nop\n\t""nop\n\t"
      "nop\n\t""nop\n\t""nop\n\t""nop\n\t"
      "nop\n\t""nop\n\t""nop\n\t""nop\n\t");
}

int32_t stepRateToQ16(float stepsPerSecond) {
  if (stepsPerSecond > MAX_TIMER_STEP_RATE) stepsPerSecond = MAX_TIMER_STEP_RATE;
  if (stepsPerSecond < -MAX_TIMER_STEP_RATE) stepsPerSecond = -MAX_TIMER_STEP_RATE;
  return (int32_t)(stepsPerSecond * ((float)STEP_RATE_SCALE / (float)STEPPER_TIMER_HZ));
}

void setTimerStepRate(uint8_t id, float stepsPerSecond, bool allowPastMin) {
  int32_t rateQ16 = stepRateToQ16(stepsPerSecond);
  uint8_t sreg = SREG;
  noInterrupts();
  gauges[id].timerStepRateQ16 = rateQ16;
  gauges[id].timerAllowPastMin = allowPastMin;
  if (rateQ16 == 0) {
    gauges[id].timerStepAccumulatorQ16 = 0;
  }
  SREG = sreg;
}

void stopTimerStepping(uint8_t id) {
  uint8_t sreg = SREG;
  noInterrupts();
  gauges[id].timerStepRateQ16 = 0;
  gauges[id].timerStepAccumulatorQ16 = 0;
  gauges[id].timerAllowPastMin = false;
  SREG = sreg;
}

void configureStepperHardware(uint8_t id) {
  stepperHardware[id].stepPort = portOutputRegister(digitalPinToPort(gaugeConfigs[id].stepPin));
  stepperHardware[id].stepMask = digitalPinToBitMask(gaugeConfigs[id].stepPin);
  stepperHardware[id].dirPort  = portOutputRegister(digitalPinToPort(gaugeConfigs[id].dirPin));
  stepperHardware[id].dirMask  = digitalPinToBitMask(gaugeConfigs[id].dirPin);
}

void beginStepperTimer() {
  uint8_t sreg = SREG;
  noInterrupts();
  TCCR1A = 0;
  TCCR1B = 0;
  TIMSK1 = 0;
  TCNT1 = 0;
  OCR1A = (uint16_t)((F_CPU / 8UL / STEPPER_TIMER_HZ) - 1UL);
  TIFR1 = _BV(OCF1A);
  TCCR1B |= _BV(WGM12);
  TCCR1B |= _BV(CS11);
  TIMSK1 = _BV(OCIE1A);
  SREG = sreg;
}

ISR(TIMER1_COMPA_vect) {
  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    if (gauges[i].timerPulseActive) {
      writeStepPin(i, false);
      gauges[i].timerPulseActive = false;
    }
  }

  for (uint8_t i = 0; i < GAUGE_COUNT; i++) {
    Gauge& g = gauges[i];
    int32_t rateQ16 = g.timerStepRateQ16;
    if (!g.enabled || rateQ16 == 0) continue;

    int32_t incrementQ16 = rateQ16 > 0 ? rateQ16 : -rateQ16;
    int32_t accumulatorQ16 = g.timerStepAccumulatorQ16 + incrementQ16;
    if (accumulatorQ16 < STEP_RATE_SCALE) {
      g.timerStepAccumulatorQ16 = accumulatorQ16;
      continue;
    }
    g.timerStepAccumulatorQ16 = accumulatorQ16 - STEP_RATE_SCALE;

    if (rateQ16 > 0) {
      if (g.currentPos >= g.targetPos || g.currentPos >= g.maxPos) {
        g.timerStepRateQ16 = 0;
        g.timerStepAccumulatorQ16 = 0;
        continue;
      }

      writeDirectionPin(i, true);
      stepDirectionSetupDelay();
      writeStepPin(i, true);
      g.timerPulseActive = true;
      g.currentPos++;
    } else {
      if (!g.timerAllowPastMin && (g.currentPos <= g.targetPos || g.currentPos <= g.minPos)) {
        g.timerStepRateQ16 = 0;
        g.timerStepAccumulatorQ16 = 0;
        continue;
      }

      if (g.timerAllowPastMin && g.homingStepsRemaining <= 0) {
        g.timerStepRateQ16 = 0;
        g.timerStepAccumulatorQ16 = 0;
        continue;
      }

      writeDirectionPin(i, false);
      stepDirectionSetupDelay();
      writeStepPin(i, true);
      g.timerPulseActive = true;
      g.currentPos--;

      if (g.timerAllowPastMin) {
        g.homingStepsRemaining--;
        if (g.homingStepsRemaining <= 0) {
          g.timerStepRateQ16 = 0;
          g.timerStepAccumulatorQ16 = 0;
        }
      }
    }
  }
}

// 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 = gaugeConfigs[id].enablePin;
  if (pin < 0) return;

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

// 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];
  stopTimerStepping(id);
  g.homingState = HS_START;
  g.homed = false;
  g.velocity = 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 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;
      atomicWriteLong(g.homingStepsRemaining, g.homingBackoffSteps);
      setTimerStepRate(id, -g.homingSpeed, true);
      g.homingState = HS_BACKING;
      break;

    case HS_BACKING:
      if (atomicReadLong(g.homingStepsRemaining) <= 0) {
        stopTimerStepping(id);
        g.homingState = HS_SETTLE;
        g.homingStateStartMs = nowMs;
      }
      break;

    case HS_SETTLE:
      if (nowMs - g.homingStateStartMs >= 100) {
        stopTimerStepping(id);
        atomicWriteLong(g.currentPos, 0);
        atomicWriteLong(g.targetPos, 0);
        g.velocity = 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;

  long currentPos = atomicReadLong(g.currentPos);
  if (g.sweepTowardMax) {
    atomicWriteLong(g.targetPos, g.maxPos);
    if (currentPos >= g.maxPos && absf(g.velocity) < 1.0f) {
      g.sweepTowardMax = false;
      atomicWriteLong(g.targetPos, g.minPos);
    }
  } else {
    atomicWriteLong(g.targetPos, g.minPos);
    if (currentPos <= g.minPos && absf(g.velocity) < 1.0f) {
      g.sweepTowardMax = true;
      atomicWriteLong(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) {
    stopTimerStepping(id);
    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) {
    stopTimerStepping(id);
    return;
  }

  long currentPos = atomicReadLong(g.currentPos);
  long targetPos = atomicReadLong(g.targetPos);
  long error = targetPos - currentPos;

  if (error == 0 && absf(g.velocity) < 0.01f) {
    g.velocity = 0.0f;
    stopTimerStepping(id);
    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;
  }

  setTimerStepRate(id, g.velocity, false);
}

// 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);
    atomicWriteLong(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 id; float speed;
  if (sscanf(line.c_str(), "SPEED %d %f", &id, &speed) == 2) {
    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;
  }
  return false;
}

// Parses `ACCEL <id> <accel>` and updates the acceleration limit.
// Replies: `OK`, `ERR BAD_ID`, `ERR BAD_ACCEL`.
bool parseAccel(const String& line) {
  int id; float accel;
  if (sscanf(line.c_str(), "ACCEL %d %f", &id, &accel) == 2) {
    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;
  }
  return false;
}

// 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];
    stopTimerStepping(id);
    atomicWriteLong(g.currentPos, 0);
    atomicWriteLong(g.targetPos, 0);
    g.velocity = 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 id; float accel, speed;
  if (sscanf(line.c_str(), "SWEEP %d %f %f", &id, &accel, &speed) == 3) {
    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;
      stopTimerStepping(id);
      sendReply("OK");
      return true;
    }
    g.accel = accel;
    g.maxSpeed = speed;
    g.sweepEnabled = true;
    g.sweepTowardMax = true;
    atomicWriteLong(g.targetPos, g.maxPos);
    sendReply("OK");
    return true;
  }
  return false;
}

// 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(atomicReadLong(gauges[i].currentPos));
      CMD_PORT.print(' ');
      CMD_PORT.print(atomicReadLong(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(gaugeConfigs[i].dirPin, OUTPUT);
    pinMode(gaugeConfigs[i].stepPin, OUTPUT);
    configureStepperHardware(i);

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

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

    gauges[i].minPos             = gaugeConfigs[i].minPos;
    gauges[i].maxPos             = gaugeConfigs[i].maxPos;
    gauges[i].homingBackoffSteps = gaugeConfigs[i].homingBackoffSteps;
    gauges[i].maxSpeed           = (float)gaugeConfigs[i].maxSpeed;
    gauges[i].accel              = (float)gaugeConfigs[i].accel;
    gauges[i].homingSpeed        = (float)gaugeConfigs[i].homingSpeed;

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

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