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


#define i2c_Address 0x3c //initialize with the I2C addr 0x3C Typically eBay OLED's
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SH1106G display = Adafruit_SH1106G(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[] = {0xE0, 0x5A, 0x1B, 0x5F, 0x14, 0xEC};
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 OnDataSent(const uint8_t *mac_addr, esp_now_send_status_t status) {
  Serial.print("\r\nLast data transmission status:\t");
  Serial.println(status == ESP_NOW_SEND_SUCCESS ? "Success" : "Failure");
}

void sendData(int servoNumber, int angle) {
  // Prepare data packet
  ReceivedData data;
  data.servoNumber = servoNumber;
  data.angle = angle;
  Serial.print("Servo Number: ");
  Serial.print(servoNumber);
  Serial.print("\t");
  Serial.print("Servo Angle: ");
  Serial.println(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");
  }
  delay(100);
}

void handleButtonPress() {
  if (digitalRead(BUTTON_1) == LOW)
  {
    sendData(0, 180); // Move all servos to 90 degrees
    sendData(1, 180); // Move all servos to 90 degrees
    sendData(2, 180); // Move all servos to 90 degrees
    sendData(3, 180); // Move all servos to 90 degrees
    sendData(4, 180); // 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;
    delay(500);
  }
  else if (digitalRead(BUTTON_2) == LOW)
  {
    sendData(0, 0); // Move all servos to 0 degrees
    sendData(1, 0); // Move all servos to 0 degrees
    sendData(2, 0); // Move all servos to 0 degrees
    sendData(3, 0); // Move all servos to 0 degrees
    sendData(4, 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;
    delay(500);
  }
  else if (digitalRead(BUTTON_3) == LOW)
  {
    sendData(0, 0); // Move servo 1 to 0 degrees
    sendData(3, 0); // Move servo 4 to 0 degrees
    sendData(4, 0); // Move servo 5 to 0 degrees
    sendData(1, 180); // Move servo 2 to 180 degrees
    sendData(2, 180); // Move servo 3 to 180 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;
    delay(500);
  }
  else if (digitalRead(BUTTON_4) == LOW)
  {
    currentMode = MODE_4;
    flag = false;
    delay(500);
  }
}

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

  // 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
  display.begin(i2c_Address, true); // Address 0x3C default
  display.display();
  delay(2000);
  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;
  }

  // After initializing ESP-NOW, register the data transmission callback
  esp_now_register_send_cb(OnDataSent);

  // 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;
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.println("> Mode 1");
  display.println("> Mode 2");
  display.println("> Mode 3");
  display.println("> Mode 4");
  display.display();
}

void calibratePotentiometers() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.println("Calibration");
  display.println("Process");
  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]);
    Serial.println(potentiometerCalibrationValues[i][0]);
  }
  Serial.println();
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.println("Calibration");
  display.println("Process");
  display.println("Closed Hand");
  display.display();
  delay(1000);
  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]);
    Serial.println(potentiometerCalibrationValues[i][1]);
  }
  delay(1000);
}

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) {
    delay(1000);
    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][0], potentiometerCalibrationValues[i][1], 0, 180);
      if (angle < 0)
      {
        sendData(i,0);
      }
      else if (angle>180)
      {
        sendData(i,180);
      }
      else
        sendData(i, angle);
    }
  }
  delay(100);
}