#include <WiFi.h>
#include <esp_now.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SH110X.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SH1106G display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define BUTTON_1 12
#define BUTTON_2 13
#define BUTTON_3 14
#define BUTTON_4 15

#define POTENTIOMETER_COUNT 5
#define SERVO_COUNT 5

// Replace the MAC Address with the one of your receiver ESP32
uint8_t broadcastAddress[] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
int PotPin[5]={32,35,34,39,36};
bool flag=false;

struct PotentiometerData {
  int value;
};

struct ReceivedData {
  int servoNumber;
  int angle;
};

enum Mode {
  MODE_1,
  MODE_2,
  MODE_3,
  MODE_4
};

Mode currentMode = MODE_1;

PotentiometerData potentiometerData[POTENTIOMETER_COUNT];
int potentiometerCalibrationValues[POTENTIOMETER_COUNT][2]; // [potentiometerIndex][0] for open hand, [potentiometerIndex][1] for closed hand

void sendData(int servoNumber, int angle) {
  // Prepare data packet
  ReceivedData data;
  data.servoNumber = servoNumber;
  data.angle = angle;

  // Send data packet
  esp_err_t result = esp_now_send(broadcastAddress, (uint8_t *) &data, sizeof(data));
  if (result == ESP_OK) {
    Serial.println("Data sent successfully");
  } else {
    Serial.println("Error sending data");
  }
}

void handleButtonPress() {
  if (digitalRead(BUTTON_1) == LOW) {
    sendData(-1, 90); // Move all servos to 90 degrees
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SH110X_WHITE);
    display.setCursor(0, 0);
    display.println("Mode 1");
    display.display();
    currentMode = MODE_1;
    flag = false;
  } else if (digitalRead(BUTTON_2) == LOW) {
    sendData(-1, 0); // Move all servos to 0 degrees
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SH110X_WHITE);
    display.setCursor(0, 0);
    display.println("Mode 2");
    display.display();
    currentMode = MODE_2;
    flag = false;
  } else if (digitalRead(BUTTON_3) == LOW) {
    sendData(0, 0); // Move servo 1,4,5 to 0 degrees
    sendData(1, 90); // Move servo 2,3 to 90 degrees
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SH110X_WHITE);
    display.setCursor(0, 0);
    display.println("Mode 3");
    display.display();
    currentMode = MODE_3;
    flag = false;
  } else if (digitalRead(BUTTON_4) == LOW) {
    currentMode = MODE_4;
    flag = false;
  }
}

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

  // Initialize buttons
  pinMode(BUTTON_1, INPUT_PULLUP);
  pinMode(BUTTON_2, INPUT_PULLUP);
  pinMode(BUTTON_3, INPUT_PULLUP);
  pinMode(BUTTON_4, INPUT_PULLUP);

  // Initialize OLED display
  if (!display.begin(SH110X_SETCOMPINS, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  display.display();
  
  // Set device as a Wi-Fi Station
  WiFi.mode(WIFI_STA);

  // Init ESP-NOW
  if (esp_now_init() != ESP_OK) {
    Serial.println("Error initializing ESP-NOW");
    return;
  }

  // Register peer
  esp_now_peer_info_t peerInfo;
  memcpy(peerInfo.peer_addr, broadcastAddress, 6);
  peerInfo.channel = 0;
  peerInfo.encrypt = false;

  if (esp_now_add_peer(&peerInfo) != ESP_OK) {
    Serial.println("Failed to add peer");
    return;
  }
}

void calibratePotentiometers() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.println("Calibration Mode");
  display.println("Open Hand");
  display.display();
  while(digitalRead(BUTTON_4) == HIGH); // Wait for user to open hand

  for (int i = 0; i < POTENTIOMETER_COUNT; i++) {
    potentiometerCalibrationValues[i][0] = analogRead(PotPin[i]);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.println("Closed Hand");
  display.display();
  while(digitalRead(BUTTON_4) == HIGH); // Wait for user to close hand

  for (int i = 0; i < POTENTIOMETER_COUNT; i++) {
    potentiometerCalibrationValues[i][1] = analogRead(PotPin[i]);
  }
}

void loop() {
  // Read potentiometer values
  for (int i = 0; i < POTENTIOMETER_COUNT; i++) {
    potentiometerData[i].value = analogRead(PotPin[i]);
  }

  // Handle button presses and send data accordingly
  handleButtonPress();

  // In Mode 4, calibrate potentiometers
  if (currentMode == MODE_4 && flag==false) {
    calibratePotentiometers();
    flag = true; // Return to Mode 1 after calibration
     display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SH110X_WHITE);
    display.setCursor(0, 0);
    display.println("Mode 4");
    display.display();
  }

  // Send potentiometer data for servo control
  if (flag) {
    for (int i = 0; i < POTENTIOMETER_COUNT; i++) {
      int angle = map(potentiometerData[i].value, potentiometerCalibrationValues[i][1], potentiometerCalibrationValues[i][0], 0, 90);
      sendData(i, angle);
    }
  }
  delay(100);
}