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

// OLED display settings
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Rotary encoder pins
#define ENCODER_PIN_A 2
#define ENCODER_PIN_B 3
#define ENCODER_SWITCH 4
Encoder myEnc(ENCODER_PIN_A, ENCODER_PIN_B);

// Z Limit pin
#define Z_LIMIT_PIN 5

// Output pins
#define OUTPUT_X1 6
#define OUTPUT_X10 7
#define OUTPUT_X100 8

// OLED display object
U8G2_SH1106_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);

// Variables
volatile bool zLimitTriggered = false;
int currentState = 0; // Default state
int lastState = 0;
int storedState = 0;
long oldPosition = -999;

// Function to handle Z Limit interrupt
void zLimitISR() {
  zLimitTriggered = true;
}

void setup() {
  // Initialize serial communication
  Serial.begin(9600);

  // Initialize OLED display
  u8g2.begin();

  // Set Z Limit pin as input with interrupt
  pinMode(Z_LIMIT_PIN, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(Z_LIMIT_PIN), zLimitISR, FALLING);

  // Set encoder switch pin as input
  pinMode(ENCODER_SWITCH, INPUT_PULLUP);

  // Set output pins as output
  pinMode(OUTPUT_X1, OUTPUT);
  pinMode(OUTPUT_X10, OUTPUT);
  pinMode(OUTPUT_X100, OUTPUT);

  // Set initial state of output pins
  digitalWrite(OUTPUT_X1, HIGH);
  digitalWrite(OUTPUT_X10, HIGH);
  digitalWrite(OUTPUT_X100, HIGH);
  digitalWrite(OUTPUT_X1, LOW);
  delay(300);
  digitalWrite(OUTPUT_X1, HIGH);

  // Display initial state
  Serial.print("Initial state: ");
  Serial.println(currentState);
  displayState(storedState);
}

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

  // Update state if the encoder position changes
  if (newPosition != oldPosition) {
    oldPosition = newPosition;
    currentState = (newPosition % 3 + 3) % 3; // Ensure state is 0, 1, or 2
    Serial.print("New state: ");
    Serial.println(currentState);
  }

  // Check if encoder switch is pressed
  if (digitalRead(ENCODER_SWITCH) == LOW) {
    storedState = currentState; // Store the current state in the register
    Serial.print("Stored state: ");
    Serial.println(storedState);
    displayState(storedState);
    triggerOutput(storedState);
    delay(200); // Debounce delay
  }

  // Handle Z Limit interrupt
  if (zLimitTriggered) {
    Serial.println("Z Limit triggered!");
    if (storedState == 2) {
      digitalWrite(OUTPUT_X10, LOW);
      delay(300);
      digitalWrite(OUTPUT_X10, HIGH);
    }
    zLimitTriggered = false;
    // Add your additional Z Limit handling code here
  }
}

// Function to display the current state on the OLED
void displayState(int state) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB14_tr);
  u8g2.setCursor(0, 30);
  switch (state) {
    case 0:
      u8g2.print("x1");
      break;
    case 1:
      u8g2.print("x10");
      break;
    case 2:
      u8g2.print("x100");
      break;
    default:
      u8g2.print("Unknown");
  }
  u8g2.sendBuffer();
}

// Function to trigger output pins based on the stored state
void triggerOutput(int state) {
  switch (state) {
    case 0:
      digitalWrite(OUTPUT_X1, LOW);
      delay(300);
      digitalWrite(OUTPUT_X1, HIGH);
      break;
    case 1:
      digitalWrite(OUTPUT_X10, LOW);
      delay(300);
      digitalWrite(OUTPUT_X10, HIGH);
      break;
    case 2:
      digitalWrite(OUTPUT_X100, LOW);
      delay(300);
      digitalWrite(OUTPUT_X100, HIGH);
      break;
  }
}