#include <Arduino.h>
#include <Servo.h>
#include <Wire.h> // Include the Wire library for I2C communication
#include <Adafruit_GFX.h> // Include the Adafruit graphics library
#include <Adafruit_SSD1306.h> // Include the Adafruit SSD1306 OLED display library

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)

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

const int buttonPins[] = {22, 23}; // Button pins for displaying 1 and 2
const int numButtons = sizeof(buttonPins) / sizeof(buttonPins[0]);
const int servoPins[] = {2, 3}; // Servo pins
const int numServos = sizeof(servoPins) / sizeof(servoPins[0]);

Servo servos[numServos]; // Create servo objects

void setup() {
  // Initialize OLED display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x64
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  display.clearDisplay(); // Clear the display buffer
  display.setTextColor(SSD1306_WHITE); // Set text color to white

  // Display "0" on OLED
  display.setTextSize(2);
  display.setCursor(0, 0);
  display.print("0");
  display.display();
  delay(2000);
  display.clearDisplay();

  // Initialize buttons
  for (int i = 0; i < numButtons; i++) {
    pinMode(buttonPins[i], INPUT_PULLUP); // Enable internal pull-up resistor for each button pin
  }

  // Attach all servo pins
  for (int i = 0; i < numServos; i++) {
    servos[i].attach(servoPins[i]);
    servos[i].write(0); // Move servos to 0 degrees
    delay(20); // Wait for the servo to reach the position
  }

  // Display initial message on OLED
  display.setTextSize(1);
  display.setCursor(0,0);
  display.println(F("Press buttons to"));
  display.println(F("control servos"));
  display.display();
  delay(2000);
  display.clearDisplay();
}

void loop() {
  // Check each button for a press
  for (int i = 0; i < numButtons; i++) {
    if (digitalRead(buttonPins[i]) == LOW) {
      // Display the corresponding number on OLED
      display.clearDisplay();
      display.setTextSize(2);
      display.setCursor(0, 0);
      display.print(i + 1); // Display 1 for button 22, 2 for button 23
      display.display();
      delay(2000); // Display for 2 seconds
      display.clearDisplay();
    }
  }
}