/*
  JAXLAB servo tester, for an ESP32-C3 SuperMini, a 128 x 64 SSD1306 OLED and a KY-040 rotary encoder.

  SELECT screen   Turn to move the cursor. Press to tick a servo, or ALL to tick or clear all six.
                  Move to GO and press to continue.
  MANUAL screen   Turn clockwise for forward, anticlockwise for backward.
                  Press for SWING. Hold to return to SELECT.
  SWING screen    The targeted servos sweep back and forth by themselves. Turn to change the speed.
                  Press for MANUAL. Hold to return to SELECT.

  Arduino IDE
    Board:      ESP32C3 Dev Module (esp32 by Espressif Systems)
    Settings:   USB CDC On Boot: Enabled
    Libraries:  ESP32Servo, Adafruit SSD1306, Adafruit GFX Library

  All the values you're most likely to change are in the settings below.
*/

#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ESP32Servo.h>

// ---------------------------------------------------------------------------
// Settings
// ---------------------------------------------------------------------------
const uint8_t SERVO_COUNT = 6;
const uint8_t SERVO_PINS[SERVO_COUNT] = {0, 1, 3, 4, 5, 6};  // header channels 1 to 6
const uint8_t PIN_SDA = 8;
const uint8_t PIN_SCL = 9;
const uint8_t PIN_ENC_CLK = 10;
const uint8_t PIN_ENC_DT = 20;
const uint8_t PIN_ENC_SW = 21;
const uint8_t OLED_ADDRESS = 0x3C;

const int ANGLE_MIN = 10;             // kept short of 0 and 180, so servos don't strain against their end stops
const int ANGLE_MAX = 170;
const int ANGLE_START = 90;           // where a servo goes when it's first targeted
const int MANUAL_STEP = 2;            // degrees for each click of the encoder
const int PULSE_MIN_US = 500;         // pulse width at 0 degrees
const int PULSE_MAX_US = 2400;        // pulse width at 180 degrees

const int SWING_SPEED_START = 60;     // degrees per second
const int SWING_SPEED_MIN = 10;
const int SWING_SPEED_MAX = 180;
const int SWING_SPEED_STEP = 5;       // for each click of the encoder

const unsigned long ATTACH_GAP_MS = 200;  // servos start one at a time, to limit the surge on the 2 A supply
const unsigned long SWING_GAP_MS = 250;   // and start swinging one after another, so they don't all peak together
const unsigned long SWING_TICK_MS = 20;
const unsigned long DEBOUNCE_MS = 25;
const unsigned long HOLD_MS = 700;
const unsigned long REDRAW_MS = 200;

const int ENC_COUNTS_PER_CLICK = 4;   // try 2 if one click moves two steps
const bool ENC_REVERSE = false;       // set to true if clockwise goes backward

// ---------------------------------------------------------------------------
// State
// ---------------------------------------------------------------------------
enum Screen { SCREEN_SELECT, SCREEN_MANUAL, SCREEN_SWING };
enum ButtonEvent { BUTTON_NONE, BUTTON_PRESS, BUTTON_HOLD };

const uint8_t CURSOR_ALL = SERVO_COUNT;
const uint8_t CURSOR_GO = SERVO_COUNT + 1;
const uint8_t CURSOR_COUNT = SERVO_COUNT + 2;
const uint8_t ALL_MASK = (1 << SERVO_COUNT) - 1;

Adafruit_SSD1306 display(128, 64, &Wire, -1);
Servo servos[SERVO_COUNT];

Screen screen = SCREEN_SELECT;
uint8_t selectedMask = 0;
uint8_t cursorPos = 0;
bool attachedFlags[SERVO_COUNT] = {};
float angles[SERVO_COUNT];
int8_t directions[SERVO_COUNT];
unsigned long swingStarts[SERVO_COUNT];
int swingSpeed = SWING_SPEED_START;
unsigned long lastSwingTick = 0;
unsigned long lastRedraw = 0;
bool redrawNeeded = true;

volatile int32_t encoderCounts = 0;
volatile uint8_t encoderState = 0;
int32_t encoderUsed = 0;

bool buttonRaw = false;
bool buttonStable = false;
bool holdFired = false;
unsigned long buttonChanged = 0;
unsigned long buttonPressed = 0;

// ---------------------------------------------------------------------------
// Encoder and button
// ---------------------------------------------------------------------------
void IRAM_ATTR onEncoderChange() {
  static const int8_t STEPS[16] = {0, -1, 1, 0, 1, 0, 0, -1, -1, 0, 0, 1, 0, 1, -1, 0};
  uint8_t pins = (digitalRead(PIN_ENC_CLK) << 1) | digitalRead(PIN_ENC_DT);
  encoderState = ((encoderState << 2) | pins) & 0x0F;
  encoderCounts = encoderCounts + STEPS[encoderState];
}

// Returns the whole clicks turned since the last call, positive for clockwise.
int readEncoderClicks() {
  noInterrupts();
  int32_t counts = encoderCounts;
  interrupts();
  int clicks = (counts - encoderUsed) / ENC_COUNTS_PER_CLICK;
  encoderUsed += clicks * ENC_COUNTS_PER_CLICK;
  return ENC_REVERSE ? -clicks : clicks;
}

ButtonEvent readButton() {
  unsigned long now = millis();
  bool raw = digitalRead(PIN_ENC_SW) == LOW;
  if (raw != buttonRaw) {
    buttonRaw = raw;
    buttonChanged = now;
  }
  if (now - buttonChanged >= DEBOUNCE_MS && buttonRaw != buttonStable) {
    buttonStable = buttonRaw;
    if (buttonStable) {
      buttonPressed = now;
      holdFired = false;
    } else if (!holdFired) {
      return BUTTON_PRESS;
    }
  }
  if (buttonStable && !holdFired && now - buttonPressed >= HOLD_MS) {
    holdFired = true;
    return BUTTON_HOLD;
  }
  return BUTTON_NONE;
}

// ---------------------------------------------------------------------------
// Servos
// ---------------------------------------------------------------------------
bool isSelected(uint8_t index) {
  return selectedMask & (1 << index);
}

void writeServo(uint8_t index) {
  servos[index].write((int)(angles[index] + 0.5f));
}

// Starts the targeted servos one at a time, and releases the rest so they draw no holding current.
void applySelection() {
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    if (attachedFlags[i] && !isSelected(i)) {
      servos[i].detach();
      attachedFlags[i] = false;
    }
  }
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    if (!attachedFlags[i] && isSelected(i)) {
      servos[i].setPeriodHertz(50);
      servos[i].attach(SERVO_PINS[i], PULSE_MIN_US, PULSE_MAX_US);
      writeServo(i);
      attachedFlags[i] = true;
      delay(ATTACH_GAP_MS);
    }
  }
}

void startSwing() {
  unsigned long now = millis();
  unsigned long start = now;
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    if (isSelected(i)) {
      swingStarts[i] = start;
      directions[i] = angles[i] >= ANGLE_MAX ? -1 : 1;
      start += SWING_GAP_MS;
    }
  }
  lastSwingTick = now;
}

void updateSwing() {
  unsigned long now = millis();
  if (now - lastSwingTick < SWING_TICK_MS) {
    return;
  }
  float seconds = (now - lastSwingTick) / 1000.0f;
  lastSwingTick = now;
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    // The subtraction stays correct when millis() wraps round.
    if (!isSelected(i) || (long)(now - swingStarts[i]) < 0) {
      continue;
    }
    angles[i] += directions[i] * swingSpeed * seconds;
    if (angles[i] >= ANGLE_MAX) {
      angles[i] = ANGLE_MAX;
      directions[i] = -1;
    } else if (angles[i] <= ANGLE_MIN) {
      angles[i] = ANGLE_MIN;
      directions[i] = 1;
    }
    writeServo(i);
  }
}

void moveManual(int clicks) {
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    if (isSelected(i)) {
      angles[i] = constrain(angles[i] + clicks * MANUAL_STEP, (float)ANGLE_MIN, (float)ANGLE_MAX);
      writeServo(i);
    }
  }
}

// ---------------------------------------------------------------------------
// Display
// ---------------------------------------------------------------------------
void drawStarting() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 28);
  display.print(F("Starting servos..."));
  display.display();
}

void drawSelect() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print(F("SELECT SERVOS"));

  const int boxSize = 17;
  const int gap = 4;
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    int x = i * (boxSize + gap) + 2;
    int y = 13;
    if (isSelected(i)) {
      display.fillRect(x, y, boxSize, boxSize, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK);
    } else {
      display.drawRect(x, y, boxSize, boxSize, SSD1306_WHITE);
      display.setTextColor(SSD1306_WHITE);
    }
    display.setCursor(x + 6, y + 5);
    display.print(i + 1);
    display.setTextColor(SSD1306_WHITE);
    if (cursorPos == i) {
      display.fillRect(x, y + boxSize + 2, boxSize, 2, SSD1306_WHITE);
    }
  }

  const int rowY = 41;
  display.setCursor(14, rowY);
  display.print(selectedMask == ALL_MASK ? F("[x] ALL") : F("[ ] ALL"));
  display.setCursor(86, rowY);
  display.print(F("GO >"));
  if (cursorPos == CURSOR_ALL) {
    display.fillRect(14, rowY + 10, 42, 2, SSD1306_WHITE);
  } else if (cursorPos == CURSOR_GO) {
    display.fillRect(86, rowY + 10, 24, 2, SSD1306_WHITE);
  }

  display.setCursor(0, 56);
  display.print(cursorPos == CURSOR_GO && selectedMask == 0 ? F("Tick a servo first") : F("Press to tick"));
}

void drawAngles() {
  uint8_t shown = 0;
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    if (isSelected(i)) {
      display.setCursor(shown % 2 == 0 ? 0 : 66, 13 + (shown / 2) * 10);
      display.print('S');
      display.print(i + 1);
      display.print(F(": "));
      display.print((int)(angles[i] + 0.5f));
      shown++;
    }
  }
}

void drawManual() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print(F("MANUAL"));
  drawAngles();
  display.setCursor(0, 46);
  display.print(F("Turn: back / forward"));
  display.setCursor(0, 56);
  display.print(F("Press: swing Hold: <"));
}

void drawSwing() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print(F("AUTO SWING"));
  drawAngles();
  display.setCursor(0, 46);
  display.print(F("Turn: speed "));
  display.print(swingSpeed);
  display.print(F(" d/s"));
  display.setCursor(0, 56);
  display.print(F("Press: manual Hold: <"));
}

void drawScreen() {
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  switch (screen) {
    case SCREEN_SELECT: drawSelect(); break;
    case SCREEN_MANUAL: drawManual(); break;
    case SCREEN_SWING: drawSwing(); break;
  }
  display.display();
  lastRedraw = millis();
  redrawNeeded = false;
}

// ---------------------------------------------------------------------------
// Screens
// ---------------------------------------------------------------------------
void handleSelect(int clicks, ButtonEvent button) {
  if (clicks != 0) {
    int moved = (cursorPos + clicks) % CURSOR_COUNT;
    cursorPos = moved < 0 ? moved + CURSOR_COUNT : moved;
    redrawNeeded = true;
  }
  if (button != BUTTON_PRESS) {
    return;
  }
  redrawNeeded = true;
  if (cursorPos < SERVO_COUNT) {
    selectedMask ^= (1 << cursorPos);
  } else if (cursorPos == CURSOR_ALL) {
    selectedMask = selectedMask == ALL_MASK ? 0 : ALL_MASK;
  } else if (selectedMask != 0) {
    drawStarting();
    applySelection();
    screen = SCREEN_MANUAL;
  }
}

void handleManual(int clicks, ButtonEvent button) {
  if (clicks != 0) {
    moveManual(clicks);
    redrawNeeded = true;
  }
  if (button == BUTTON_PRESS) {
    startSwing();
    screen = SCREEN_SWING;
    redrawNeeded = true;
  } else if (button == BUTTON_HOLD) {
    screen = SCREEN_SELECT;
    redrawNeeded = true;
  }
}

void handleSwing(int clicks, ButtonEvent button) {
  updateSwing();
  if (clicks != 0) {
    swingSpeed = constrain(swingSpeed + clicks * SWING_SPEED_STEP, SWING_SPEED_MIN, SWING_SPEED_MAX);
    redrawNeeded = true;
  }
  if (button == BUTTON_PRESS) {
    screen = SCREEN_MANUAL;
    redrawNeeded = true;
  } else if (button == BUTTON_HOLD) {
    screen = SCREEN_SELECT;
    redrawNeeded = true;
  }
  if (millis() - lastRedraw >= REDRAW_MS) {
    redrawNeeded = true;  // keeps the angles current while they swing
  }
}

// ---------------------------------------------------------------------------
void setup() {
  pinMode(PIN_ENC_CLK, INPUT_PULLUP);
  pinMode(PIN_ENC_DT, INPUT_PULLUP);
  pinMode(PIN_ENC_SW, INPUT_PULLUP);

  encoderState = (digitalRead(PIN_ENC_CLK) << 1) | digitalRead(PIN_ENC_DT);
  attachInterrupt(digitalPinToInterrupt(PIN_ENC_CLK), onEncoderChange, CHANGE);
  attachInterrupt(digitalPinToInterrupt(PIN_ENC_DT), onEncoderChange, CHANGE);

  // Nothing attaches until you press GO, so the servos stay still at power-on.
  for (uint8_t i = 0; i < SERVO_COUNT; i++) {
    angles[i] = ANGLE_START;
    directions[i] = 1;
  }
  for (uint8_t timer = 0; timer < 4; timer++) {
    ESP32PWM::allocateTimer(timer);
  }

  Wire.begin(PIN_SDA, PIN_SCL);
  display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDRESS);
  drawScreen();
}

void loop() {
  int clicks = readEncoderClicks();
  ButtonEvent button = readButton();
  switch (screen) {
    case SCREEN_SELECT: handleSelect(clicks, button); break;
    case SCREEN_MANUAL: handleManual(clicks, button); break;
    case SCREEN_SWING: handleSwing(clicks, button); break;
  }
  if (redrawNeeded) {
    drawScreen();
  }
}