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

// Define the OLED display size and I2C address
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Define the step and direction pins for the stepper motor
const int stepPin = 0;
const int dirPin = 2;

// Define the button pins
const int buttonNext = 4;
const int buttonSelect = 5;

// Define the potentiometer pin
const int potPin = A0;

// Define the RGB LED pins
const int redPin = 14;
const int greenPin = 15;
const int bluePin = 16;

// Menu variables
int menuIndex = 0;
bool loadMenu = false;
bool forward = true;  // True for clockwise, False for counterclockwise
bool motorRunning = false;
bool motorRunningInSequence = false;
unsigned long lastStepTime = 0;
int stepDelay = 1000; // Initial delay between steps in microseconds
int stepsToRun = 400; // Number of steps to run before changing direction

void setup() {
  Serial.begin(9600);
  
  // Initialize the OLED display
if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
   }
  //display.display();  // Update the display
  display.clearDisplay();
   // Display the startup message
  display.setTextSize(2);  // Set text size
  display.setTextColor(SSD1306_WHITE);  // Set text color
  display.setCursor(0, 0);  // Set cursor position
  display.println(F("A project by"));  // Print the first line
  display.println(F("Kishore & Son's"));  // Print the second line
  display.display();  // Update the display
  delay(3000);  // Keep the message on screen for 2 seconds

  // Clear the display again to prepare for the main menu
  display.clearDisplay();


  // Set the pins as outputs
  pinMode(stepPin, OUTPUT);
  pinMode(dirPin, OUTPUT);

  // Set the buttons as inputs
  pinMode(buttonNext, INPUT_PULLUP);
  pinMode(buttonSelect, INPUT_PULLUP);

  // Initialize the RGB LED pins
  pinMode(redPin, OUTPUT);
  pinMode(greenPin, OUTPUT);
  pinMode(bluePin, OUTPUT);
  setRGBColor(255, 0, 0); // Start with the LED red (idle state)

  // Show the initial menu
  showMenu();
}

void loop() {
  // Check for button presses
  if (digitalRead(buttonNext) == LOW) {
    delay(200); // Debounce delay
    if (!motorRunning && !motorRunningInSequence) {
      menuIndex = (menuIndex + 1) % (loadMenu ? 2 : 2);
      showMenu();
    }
  }

  if (digitalRead(buttonSelect) == LOW) {
    delay(200); // Debounce delay
    if (motorRunning) {
      motorRunning = false;
      setRGBColor(255, 0, 0); // Set color to red when idle
      showMenu();
    } else if (motorRunningInSequence) {
      motorRunningInSequence = false;
      setRGBColor(255, 0, 0); // Set color to red when idle
      showMenu();
    } else {
      if (!loadMenu && menuIndex == 0) {
        loadMenu = true;
        menuIndex = 0;
        showMenu();
      } else if (!loadMenu && menuIndex == 1) {
        motorRunningInSequence = true;
        setRGBColor(0, 0, 255); // Set color to blue when starting motor
      } else if (loadMenu) {
        forward = (menuIndex == 1); // Reverse the logic here: Forward is 1, Reverse is 0
        loadMenu = false;
        motorRunning = true;
        setRGBColor(0, 255, 0); // Set color to green when running
      }
      showMenu();
    }
  }

  // Update RGB LED color based on motor state
  if (!motorRunning && !motorRunningInSequence) {
    setRGBColor(255, 0, 0); // Red when motor is not active
  }

  if (motorRunning) {
    runMotor();
  } else if (motorRunningInSequence) {
    startMotor();
  }
}

void showMenu() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);

  if (!loadMenu) {
    display.setCursor(0, 0);
    display.println(F("1. Load"));
    display.setCursor(0, 32);
    display.println(F("2. Start"));
    display.setCursor(0, menuIndex * 32);
  } else {
    display.setCursor(0, 0);
    display.println(F("1. Reverse")); // Changed to Reverse
    display.setCursor(0, 32);
    display.println(F("2. Forward")); // Changed to Forward
    display.setCursor(0, menuIndex * 32);
  }

  display.fillRect(0, menuIndex * 32, 128, 32, SSD1306_WHITE);
  display.setTextColor(SSD1306_BLACK);
  display.setCursor(0, menuIndex * 32);
  display.println(loadMenu ? (menuIndex == 0 ? F("1. Reverse") : F("2. Forward")) : (menuIndex == 0 ? F("1. Load") : F("2. Start")));
  display.setTextColor(SSD1306_WHITE);

  display.display();
}

void setRGBColor(int red, int green, int blue) {
  analogWrite(redPin, red);
  analogWrite(greenPin, green);
  analogWrite(bluePin, blue);
}

void runMotor() {
  // Read the potentiometer value and map it inversely to a delay range
  int potValue = analogRead(potPin);
  stepDelay = map(potValue, 0, 4095, 2000, 500); // Adjust range as needed (higher pot value = slower motor speed)

  // Set the direction based on the forward variable
  digitalWrite(dirPin, forward ? LOW : HIGH); // Swapped the HIGH/LOW logic to match the new forward/reverse settings

  // Rotate the motor
  if (micros() - lastStepTime >= stepDelay) {
    digitalWrite(stepPin, HIGH);
    delayMicroseconds(1); // Short pulse width
    digitalWrite(stepPin, LOW);
    lastStepTime = micros();
  }
}

void startMotor() {
  static int stepsCount = 0;
  static bool direction = forward; // Initial direction based on the load menu selection

  // Read the potentiometer value and map it inversely to a delay range
  int potValue = analogRead(potPin);
  stepDelay = map(potValue, 0, 4095, 2000, 500); // Adjust range as needed (higher pot value = slower motor speed)

  // Set the direction based on the direction variable
  digitalWrite(dirPin, direction ? LOW : HIGH); // Swapped the HIGH/LOW logic to match the new forward/reverse settings

  // Rotate the motor
  if (micros() - lastStepTime >= stepDelay) {
    digitalWrite(stepPin, HIGH);
    delayMicroseconds(1); // Short pulse width
    digitalWrite(stepPin, LOW);
    lastStepTime = micros();

    stepsCount++;
    if (stepsCount >= stepsToRun) {
      stepsCount = 0;
      direction = !direction; // Change direction after a set number of steps
      digitalWrite(dirPin, direction ? LOW : HIGH); // Set new direction
      delay(100); // Short delay after changing direction
    }
  }
}