#include <Adafruit_SSD1306.h>
#include <Encoder.h>
#include <AccelStepper.h>

#define OLED_RESET 4
Adafruit_SSD1306 display(OLED_RESET);

Encoder myEncoder(2, 3); // Encoder pins

const int MENU_COUNT = 2;
String menu[MENU_COUNT] = {"Extrusion", "Cleaning"};
int menuIndex = 0;
bool extrusionSelected = false;
bool optionSelected = false;

// Stepper motor configuration
#define STEP_PIN 7
#define DIR_PIN 6
#define ENABLE_PIN A3

AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

void setup() {
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C); // OLED Address
  display.clearDisplay();
  display.setTextColor(WHITE);
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("MENU:");
  display.display();
  updateDisplay();
  
  // Initialize stepper motor
  stepper.setEnablePin(ENABLE_PIN);
  stepper.setPinsInverted(false, false, true);
  stepper.enableOutputs();
}

void loop() {
  if (!extrusionSelected) {
    int menuSelection = readEncoder();

    if (menuSelection != menuIndex) {
      menuIndex = menuSelection;
      updateDisplay();
    }

    if (digitalRead(5) == HIGH) { // Button press
      if (menuIndex == 0) {
        // Code for Extrusion menu item selected
        display.clearDisplay();
        display.setCursor((display.width() - 29) / 2, 0); // Center align "WARNING"
        display.println("WARNING:");
        display.setCursor(0, 10); // Set cursor to next line
        display.println("Make sure the barrel is filled");
        display.println("before starting process");
        display.display();
        extrusionSelected = true;
      } else if (menuIndex == 1) {
        // Code for Cleaning menu item selected
        // ...
      }
    }
  } else {
    // Extrusion menu item selected
    if (!optionSelected) {
      int optionSelection = readEncoder();
      if (optionSelection == 0) {
        // Proceed option selected
        display.clearDisplay();
        display.setCursor((display.width() - 11) / 2, 0); // Center align "PROCEED"
        display.println("PROCEED");
        display.display();
      } else if (optionSelection == 1) {
        // Cancel option selected
        extrusionSelected = false;
        optionSelected = false;
        menuIndex = 0;
        updateDisplay();
      }
      
      if (digitalRead(5) == LOW) { // Button press
        if (optionSelection == 0) {
          // Code for starting the stepper motor
          stepper.move(200); // Spin the stepper motor 200 steps
          stepper.setSpeed(1000); // Set the speed of the stepper motor
          stepper.setAcceleration(500); // Set the acceleration of the stepper motor
          stepper.enableOutputs();
          while (stepper.distanceToGo() != 0) {
            stepper.run();
          }
          
          // Display completion message
          display.clearDisplay();
          display.setCursor((display.width() - 10) / 2, 0); // Center align "COMPLETE"
          display.println("COMPLETE");
          display.display();
          
          // Reset flags and menu selection
          extrusionSelected = false;
          optionSelected = false;
          menuIndex = 0;
          
          // Wait for a brief moment before returning to the main menu
          delay(2000);
          updateDisplay();
        } else if (optionSelection == 1) {
          // Cancel option selected
          extrusionSelected = false;
          optionSelected = false;
          menuIndex = 0;
          updateDisplay();
        }
      }
    }
  }
}

int readEncoder() {
  static int previousMenuSelection = 0;
  int menuSelection = previousMenuSelection + myEncoder.read() / 4;
  menuSelection = constrain(menuSelection, 0, MENU_COUNT - 1);
  previousMenuSelection = menuSelection;
  return menuSelection;
}

void updateDisplay() {
  display.clearDisplay();
  display.setCursor((display.width() - 6 * 7) / 2, 0); // Center align "MENU"
  display.println("MENU:");

  for (int i = 0; i < MENU_COUNT; i++) {
    display.setCursor(0, i == 0 ? display.getCursorY() + 8 : display.getCursorY()); // Move cursor to next line
    display.print(i == menuIndex ? "> " : "  ");
    display.print(menu[i]);
    display.println();
  }
  display.display();
}