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

// OLED display settings
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
U8G2_SH1107_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0); // Use SH1107 display

// Rotary encoder pins
#define ENCODER_PIN_A 2
#define ENCODER_PIN_B 3
#define ENCODER_SWITCH 4

// Z minus limit input pin
#define Z_MINUS_LIMIT_PIN 9

// BOB Connected input pin
#define BOB_CONNECTED_PIN 10

Encoder myEnc(ENCODER_PIN_A, ENCODER_PIN_B);

// Output pins
#define OUTPUT1 5
#define OUTPUT2 6
#define OUTPUT3 7

// MPG status input pin
#define MPG_STATUS_PIN 8

// Variables
int currentOutput = 0;
int lastOutput = 0;
long oldPosition = -999;
unsigned long lastSwitchPressTime = 0;
bool switchPressed = false;
bool lastMpgStatus = HIGH;
unsigned long lastDisplayUpdateTime = 0;
const unsigned long DISPLAY_UPDATE_INTERVAL = 100; // Update display every 100 ms

// Pulse counting
int pulseCount = 0; // Counter for pulses
int selectedState = 0; // Counter for selected state changes

void setup() {
  Serial.begin(9600);

  // Initialize OLED display
  u8g2.begin();
  u8g2.clearBuffer();
  
  // Set output pins
  pinMode(OUTPUT1, OUTPUT);
  pinMode(OUTPUT2, OUTPUT);
  pinMode(OUTPUT3, OUTPUT);
  
  // Set encoder switch pin
  pinMode(ENCODER_SWITCH, INPUT_PULLUP);
  
  // Set MPG status pin
  pinMode(MPG_STATUS_PIN, INPUT_PULLUP);

  // Set Z minus limit pin
  pinMode(Z_MINUS_LIMIT_PIN, INPUT_PULLUP);
  
  // Set BOB Connected pin
  pinMode(BOB_CONNECTED_PIN, INPUT_PULLUP);

  // Set initial state of output pins
  resetOutputs();
  
  // Display initial output
  displayOutput(currentOutput);
}

void loop() {
  long newPosition = myEnc.read() / 4;

  // Update position and pulse counting
  if (newPosition != oldPosition) {
    oldPosition = newPosition;
    // Increment or decrement pulse count based on direction
    pulseCount += (newPosition > oldPosition) ? 1 : -1;

    // Change selection only after 3 pulses
    if (abs(pulseCount) >= 3) {
      currentOutput = (currentOutput + (pulseCount > 0 ? 1 : -1)) % 3; // Ensure non-negative output
      if (currentOutput < 0) {
        currentOutput += 3; // Wrap around
      }
      selectedState++; // Increment the selected state counter
      displaySelectedOutput(currentOutput);
      lastSwitchPressTime = millis();
      switchPressed = false;
      pulseCount = 0; // Reset pulse count after changing selection
    }
  }

  // Handle encoder switch state
  bool encoderSwitchState = digitalRead(ENCODER_SWITCH) == LOW;
  if (encoderSwitchState != switchPressed) {
    if (encoderSwitchState) {
      lastOutput = currentOutput;
      setOutputLow(currentOutput);
      displayOutput(currentOutput);
    }
    switchPressed = encoderSwitchState;
  } else if (switchPressed && (millis() - lastSwitchPressTime > 1000)) {
    displayOutput(lastOutput);
    switchPressed = false;
  }

  // Check Z minus limit input
  handleZMinusLimit();

  // Check BOB Connected input
  handleBobConnected();

  // Check MPG status pin state and update display based on frequency
  unsigned long currentTime = millis();
  if (currentTime - lastDisplayUpdateTime >= DISPLAY_UPDATE_INTERVAL) {
    bool currentMpgStatus = digitalRead(MPG_STATUS_PIN);
    if (currentMpgStatus != lastMpgStatus) {
      lastMpgStatus = currentMpgStatus;
      displayOutput(currentOutput);  // Refresh display to update MPG status
    }
    lastDisplayUpdateTime = currentTime;
  }
}

void handleZMinusLimit() {
  static bool lastZMinusState = HIGH; // Initial state of Z minus limit
  bool currentZMinusState = digitalRead(Z_MINUS_LIMIT_PIN);

  if (currentZMinusState == LOW && lastZMinusState == HIGH) {
    // If the Z minus limit goes low and the last output was Output 3
    if (lastOutput == 2) { // Output 3 corresponds to index 2
      digitalWrite(OUTPUT2, LOW);  // Momentarily set Output 2 low
      delay(300);                   // Wait for 300 ms
      digitalWrite(OUTPUT2, HIGH);  // Set Output 2 back high
    }
  } else if (currentZMinusState == HIGH && lastZMinusState == LOW) {
    // If the Z minus limit goes high
    if (lastOutput == 2) { // Output 3 corresponds to index 2
      digitalWrite(OUTPUT3, LOW);  // Momentarily set Output 3 low
      delay(300);                   // Wait for 300 ms
      digitalWrite(OUTPUT3, HIGH);  // Set Output 3 back high
    }
  }
  lastZMinusState = currentZMinusState; // Update the last state
}

void handleBobConnected() {
  static bool lastBobState = HIGH; // Initial state of BOB Connected
  bool currentBobState = digitalRead(BOB_CONNECTED_PIN);

  if (currentBobState == LOW && lastBobState == HIGH) {
    // If BOB Connected goes from HIGH to LOW
    selectedState++; // Change selected state (increment if needed)
    currentOutput = 0; // Change current output to case 0 (Output 1)
    digitalWrite(OUTPUT1, LOW);  // Momentarily set Output 1 low
    delay(300);                   // Wait for 300 ms
    digitalWrite(OUTPUT1, HIGH);  // Set Output 1 back high
    displayOutput(currentOutput); // Update display with the new current output
  }
  lastBobState = currentBobState; // Update the last state
}

void setOutputLow(int output) {
  // Turn selected output pin LOW and then HIGH with a short delay
  digitalWrite(OUTPUT1 + output, LOW);
  delay(300);
  digitalWrite(OUTPUT1 + output, HIGH);
}

void displayOutput(int output) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB08_tr);
  
  u8g2.setCursor(0, 10);
  u8g2.print(digitalRead(MPG_STATUS_PIN) == LOW ? "JOG Enabled" : "*Disabled*");
  
  u8g2.setFont(u8g2_font_ncenB10_tr);
  u8g2.setCursor(0, 30);
  u8g2.print("Output: ");
  u8g2.print(outputToString(output));

  u8g2.setCursor(0, 50);
  u8g2.print("Selected State: ");
  u8g2.print(selectedState);

  u8g2.sendBuffer();
}

void displaySelectedOutput(int output) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB08_tr);
  
  u8g2.setCursor(0, 10);
  u8g2.print(digitalRead(MPG_STATUS_PIN) == LOW ? "JOG Enabled" : "**Disabled**");
  
  u8g2.setFont(u8g2_font_ncenB10_tr);
  u8g2.setCursor(0, 30);
  u8g2.print("Select");
  
  u8g2.setCursor(0, 50);
  u8g2.print(outputToString(output));
  u8g2.setCursor(0, 60);
  u8g2.print("Selected State: ");
  u8g2.print(selectedState);

  u8g2.sendBuffer();
}

String outputToString(int output) {
  switch (output) {
    case 0: return "x 1";
    case 1: return "x 10";
    case 2: return "x 100";
    default: return "Unknown";
  }
}

void resetOutputs() {
  digitalWrite(OUTPUT1, HIGH);
  digitalWrite(OUTPUT2, HIGH);
  digitalWrite(OUTPUT3, HIGH);
  digitalWrite(OUTPUT1, LOW);
  delay(300);
  digitalWrite(OUTPUT1, HIGH);
}