#include <U8g2lib.h>
#include <Wire.h>

// ---------------- OLED ----------------
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);

// ---------------- PINS ----------------
#define PIN_STEP      2
#define PIN_DIR       3
#define PIN_EN        5
#define PIN_STARTSTOP 8
#define PIN_ENDSTOP   4
#define PIN_RETURN    12
#define PIN_LEADOUT   9
#define PIN_LEADIN    11
#define PIN_SPACING   10
#define PIN_LPI_UP    6
#define PIN_LPI_DN    7

// ---------------- PARAMETERS ----------------
int lpi = 169;                    // adjustable LPI
const float DISC_RPM = 33.333f;   // disc rotation speed
const float STEPS_PER_MM = 2889.0f;  // calibrated steps per mm

const float LEAD_MMPS = 1.0f;  // lead-in/out mm/sec
const float SPACING_MMPS = 1.0f;
const float RETURN_MMPS = 15000.0f; 

const uint32_t LEADIN_MS   = 5000;
const uint32_t SPACING_MS  = 3000;

// Sweep distances in mm and display labels
const int sweepMm[] = {90, 69, 33};
const char* sweepLabel[] = {"12\"", "10\"", "7\""};
const int sweepCount = 3;
int sweepIndex = 0;

// ---------------- STATE ----------------
volatile bool motorEnabled = false;
bool forwardMode = false;
bool leadOutActive = false;
bool leadInActive = false;
bool spacingActive = false;

float currentSpeedStepsPS = 0.0f;
uint32_t modeStartTime = 0;
unsigned long runStartTime = 0;

// Timer variables
unsigned long sweepRemainingMs = 0;
unsigned long sweepTotalMs = 0; // always recalculated from LPI

// ---------------- BUTTON STATE ----------------
const unsigned long DEBOUNCE_MS = 200;
unsigned long lastButtonTime = 0;

// ---------------- HELPERS ----------------
float lpiToStepsPerSec(float lpi, float rpm) {
  float revPerSec = rpm / 60.0f;
  float linearInPerSec = revPerSec / lpi;
  float mmPerSec = linearInPerSec * 25.4f;
  return mmPerSec * STEPS_PER_MM;
}

float mmpsToStepsps(float mmps) {
  return mmps * STEPS_PER_MM;
}

void stopTimer() {
  TIMSK1 &= ~(1 << OCIE1A);
  TCCR1B = 0;
  TCCR1A = 0;
  TCNT1 = 0;
  digitalWrite(PIN_STEP, LOW);
}

void startTimerForStepsPerSec(float stepsPerSec) {
  if (stepsPerSec <= 0) {
    stopTimer();
    return;
  }
  double toggleFreq = stepsPerSec * 2.0;
  unsigned long counts = (F_CPU / toggleFreq) - 1;
  if (counts > 65535) counts = 65535;

  TCCR1A = 0;
  TCCR1B = 0;
  TCNT1 = 0;
  OCR1A = (uint16_t)counts;
  TCCR1B = (1 << WGM12) | (1 << CS10);
  TIMSK1 |= (1 << OCIE1A);
}

void setMotorEnable(bool en, float speedStepsPS = 0.0f, bool dirForward = true) {
  motorEnabled = en;
  digitalWrite(PIN_EN, en ? LOW : HIGH);
  if (en) {
    digitalWrite(PIN_DIR, dirForward ? HIGH : LOW);
    currentSpeedStepsPS = speedStepsPS;
    startTimerForStepsPerSec(speedStepsPS);
  } else {
    stopTimer();
    currentSpeedStepsPS = 0.0f;
  }
}

// ISR toggles step
ISR(TIMER1_COMPA_vect) {
  if (motorEnabled) {
    digitalWrite(PIN_STEP, !digitalRead(PIN_STEP));
  } else {
    digitalWrite(PIN_STEP, LOW);
  }
}

// ---------------- SETUP ----------------
void setup() {
  pinMode(PIN_STEP, OUTPUT);
  pinMode(PIN_DIR, OUTPUT);
  pinMode(PIN_EN, OUTPUT);

  pinMode(PIN_STARTSTOP, INPUT_PULLUP);
  pinMode(PIN_ENDSTOP, INPUT_PULLUP);
  pinMode(PIN_RETURN, INPUT_PULLUP);
  pinMode(PIN_LEADOUT, INPUT_PULLUP);
  pinMode(PIN_LEADIN, INPUT_PULLUP);
  pinMode(PIN_SPACING, INPUT_PULLUP);
  pinMode(PIN_LPI_UP, INPUT_PULLUP);
  pinMode(PIN_LPI_DN, INPUT_PULLUP);

  digitalWrite(PIN_STEP, LOW);
  digitalWrite(PIN_DIR, HIGH);
  digitalWrite(PIN_EN, HIGH);

  currentSpeedStepsPS = lpiToStepsPerSec(lpi, DISC_RPM);
  u8g2.begin();
}

// ---------------- LOOP ----------------
void loop() {
  unsigned long now = millis();

  // --- Button handling with debounce ---
  bool upPressed = digitalRead(PIN_LPI_UP) == LOW;
  bool dnPressed = digitalRead(PIN_LPI_DN) == LOW;

  if (now - lastButtonTime > DEBOUNCE_MS) {
    // Adjust LPI
    if (upPressed && !dnPressed) { lpi++; lastButtonTime = now; }
    if (!upPressed && dnPressed) { lpi--; lastButtonTime = now; }

    // Cycle sweep mode (both pressed)
    if (upPressed && dnPressed) {
      sweepIndex = (sweepIndex + 1) % sweepCount;
      lastButtonTime = now;
    }
  }

  // --- Always recalculate total time from LPI ---
  float revPerSec = DISC_RPM / 60.0f;
  float inchesPerSec = revPerSec / lpi;
  float mmPerSec = inchesPerSec * 25.4f;
  sweepTotalMs = (unsigned long)((sweepMm[sweepIndex] / mmPerSec) * 1000.0f);

  // --- Motor controls and modes ---
  bool endstopActive = (digitalRead(PIN_ENDSTOP) == LOW);
  bool startPressed  = (digitalRead(PIN_STARTSTOP) == LOW);
  bool returnPressed = (digitalRead(PIN_RETURN) == LOW);
  bool leadOutPressed = (digitalRead(PIN_LEADOUT) == LOW);
  bool leadInPressed = (digitalRead(PIN_LEADIN) == LOW);
  bool spacingPressed = (digitalRead(PIN_SPACING) == LOW);

  // --- ENDSTOP stops motor (always resets timer) ---
  if (endstopActive && motorEnabled) {
    setMotorEnable(false);
    forwardMode = false;
    leadOutActive = false;
    leadInActive = false;
    spacingActive = false;
    sweepRemainingMs = sweepTotalMs; // reset to full sweep time
  }

  // Start/stop
  static bool lastStartState = HIGH;
  if (startPressed != lastStartState) {
    lastStartState = startPressed;
    if (startPressed && !motorEnabled && !endstopActive) {
      setMotorEnable(true, lpiToStepsPerSec(lpi, DISC_RPM), true);
      forwardMode = true;
      runStartTime = now;
      sweepRemainingMs = sweepTotalMs; // initialize countdown
    } else if (startPressed && forwardMode) {
      setMotorEnable(false);
      forwardMode = false;
    }
  }

  // Return jog
  if (returnPressed) {
    setMotorEnable(true, RETURN_MMPS, false);
    forwardMode = false;
    leadOutActive = false;
    leadInActive = false;
    spacingActive = false;
  } else if (!returnPressed && !forwardMode && !leadOutActive && !leadInActive && !spacingActive) {
    setMotorEnable(false);
  }

  // Lead-out
  if (leadOutPressed && motorEnabled && forwardMode && !leadOutActive) {
    setMotorEnable(true, mmpsToStepsps(LEAD_MMPS), true);
    leadOutActive = true;
  }
  if (leadOutActive && endstopActive) {
    setMotorEnable(false);
    leadOutActive = false;
    forwardMode = false;
  }

  // Lead-in
  if (leadInPressed && !motorEnabled) {
    setMotorEnable(true, mmpsToStepsps(LEAD_MMPS), true);
    leadInActive = true;
    modeStartTime = now;
  }
  if (leadInActive && (now - modeStartTime >= LEADIN_MS)) {
    setMotorEnable(true, lpiToStepsPerSec(lpi, DISC_RPM), true);
    leadInActive = false;
    forwardMode = true;
    runStartTime = now;
    sweepRemainingMs = sweepTotalMs;
  }

  // Spacing
  if (spacingPressed && motorEnabled) {
    setMotorEnable(true, mmpsToStepsps(SPACING_MMPS), true);
    spacingActive = true;
    modeStartTime = now;
  }
  if (spacingActive && (now - modeStartTime >= SPACING_MS)) {
    setMotorEnable(true, lpiToStepsPerSec(lpi, DISC_RPM), true);
    spacingActive = false;
    runStartTime = now;
  }

  // --- Countdown update ---
  if (forwardMode && sweepRemainingMs > 0) {
    unsigned long elapsed = now - runStartTime;
    if (elapsed < sweepRemainingMs) {
      sweepRemainingMs -= elapsed;
    } else {
      sweepRemainingMs = 0;
      setMotorEnable(false);
      forwardMode = false;
    }
    runStartTime = now;
  }

  // --- OLED display ---
  unsigned long displayTimeMs = forwardMode ? sweepRemainingMs : sweepTotalMs;
  int mins = displayTimeMs / 60000;
  int secs = (displayTimeMs / 1000) % 60;

  char topLine[32];
  snprintf(topLine, sizeof(topLine), "%s @ %d", sweepLabel[sweepIndex], lpi);
  int topWidth = u8g2.getStrWidth(topLine);
  int topX = 4;

  char timeLine[16];
  snprintf(timeLine, sizeof(timeLine), "%02d:%02d", mins, secs);
  u8g2.setFont(u8g2_font_logisoso32_tn);
  int timeWidth = u8g2.getStrWidth(timeLine);
  int timeX = (128 - timeWidth) / 2;

  u8g2.clearBuffer();

  // top line (sweep + lpi)
  u8g2.setFont(u8g2_font_10x20_tf);
  u8g2.setCursor(topX, 20);   // top margin
  u8g2.print(topLine);

  // time line (big and centered)
  u8g2.setFont(u8g2_font_logisoso32_tn);
  u8g2.setCursor(timeX, 63);  // lower baseline for big font
  u8g2.print(timeLine);

  u8g2.sendBuffer();
  delay(50);
}