#include <WiFi.h>
#include <esp_now.h>
#include <Wire.h>
#include <Adafruit_SSD1306.h>
#include <AceButton.h>

using namespace ace_button;

#define SCREEN_WIDTH 128  // OLED display width, in pixels
#define SCREEN_HEIGHT 32  // OLED display height, in pixels

// OLED display setup
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define LED_BUILTIN 2  // Built-in LED for slave connection status

// GPIO pins for push buttons
const int buttonPins[] = {13, 12, 14, 27};
const int numButtons = 4;
bool buttonStates[numButtons] = {false, false, false, false};

// Slave ESP32 MAC Address
uint8_t slaveMAC[] = {0x24, 0x0A, 0xC4, 0x00, 0x01, 0x10};

// Structure for ESP-NOW data
typedef struct {
  int buttonID;
  bool state;
} ButtonData;

ButtonData buttonData;

// Structure for receiving feedback
typedef struct {
  bool relayStates[4]; // State of all 4 relays
} FeedbackData;

FeedbackData feedbackData;

// Slave Connection Status
bool slaveConnected = false;

unsigned long lastFeedbackTime = 0;
const unsigned long feedbackTimeout = 4000; // 4 seconds timeout

// **Create AceButton instances for each button**
AceButton buttons[numButtons] = {
  AceButton(buttonPins[0]),
  AceButton(buttonPins[1]),
  AceButton(buttonPins[2]),
  AceButton(buttonPins[3])
};

// **ESP-NOW Send Callback**
void onDataSent(const uint8_t *mac_addr, esp_now_send_status_t status) {
  slaveConnected = (status == ESP_NOW_SEND_SUCCESS);
  digitalWrite(LED_BUILTIN, slaveConnected ? HIGH : LOW);  // Turn on/off inbuilt LED based on connection status
  
  // Print the status to Serial Monitor
  if (slaveConnected) {
    Serial.println("Slave connected successfully.");
  } else {
    Serial.println("Failed to send data to Slave.");
  }
}

// **Fixed ESP-NOW Receive Callback Signature**
void onDataRecv(const esp_now_recv_info_t *info, const uint8_t *incomingData, int len) {
  Serial.print("Received data from: ");
  for (int i = 0; i < 6; i++) {
    Serial.print(info->src_addr[i], HEX);
    Serial.print(" ");
  }
  Serial.println();
  Serial.print("Data Length: ");
  Serial.println(len);

  // Copy the feedback data
  memcpy(&feedbackData, incomingData, sizeof(feedbackData));

  // Update the last feedback time
  lastFeedbackTime = millis();

  // Print received feedback data to Serial Monitor
  Serial.print("Received feedback from Slave: ");
  for (int i = 0; i < 4; i++) {
    Serial.print(feedbackData.relayStates[i] ? "ON " : "OFF ");
  }
  Serial.println();

  // Update the OLED display
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Relay States: ");
  display.println("-----------------");
  for (int i = 0; i < 4; i++) {
    display.print(feedbackData.relayStates[i] ? "ON " : "OFF ");
  }
  display.display();
}

// **Handle Button Press Events**
void handleButtonPress(AceButton *button, uint8_t eventType, uint8_t buttonState) {
  if (eventType == AceButton::kEventReleased) {
    int buttonIndex = button->getPin();

    // Find button index in buttonPins array
    for (int i = 0; i < numButtons; i++) {
      if (buttonPins[i] == buttonIndex) {
        buttonStates[i] = !buttonStates[i];
        buttonData.buttonID = i;
        buttonData.state = buttonStates[i];

        // Send button state to Slave ESP32
        esp_now_send(slaveMAC, (uint8_t *)&buttonData, sizeof(buttonData));
        break;
      }
    }
  }
}

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

  // **OLED Initialization**
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;) ;
  }
  delay(1000);
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.clearDisplay();

  // **Initialize Buttons**
  for (int i = 0; i < numButtons; i++) {
    pinMode(buttonPins[i], INPUT_PULLUP);
    buttons[i].setEventHandler(handleButtonPress);
  }

  // **Initialize AceButton System**
  ButtonConfig::getSystemButtonConfig()->setEventHandler(handleButtonPress);

  // **Setup ESP-NOW**
  WiFi.mode(WIFI_STA);
  if (esp_now_init() != ESP_OK) {
    Serial.println("ESP-NOW init failed");
    return;
  }

  esp_now_register_send_cb(onDataSent);
  esp_now_register_recv_cb(onDataRecv);

 // On the Master
  esp_now_peer_info_t peerInfo = {};
  memcpy(peerInfo.peer_addr, slaveMAC, 6);
  peerInfo.channel = 0;  // Ensure it's the correct channel
  peerInfo.encrypt = false;
  if (esp_now_add_peer(&peerInfo) != ESP_OK) {
    Serial.println("Failed to add peer");
    return;
  }

  // **Setup Built-in LED for Slave Connection Indication**
  pinMode(LED_BUILTIN, OUTPUT);
  digitalWrite(LED_BUILTIN, LOW);
}

void loop() {
  // **Manually call check() for each AceButton instance**
  for (int i = 0; i < numButtons; i++) {
    buttons[i].check();
  }

  // **Check if feedback is not received for more than 4 seconds**
  if (millis() - lastFeedbackTime >= feedbackTimeout) {
    digitalWrite(LED_BUILTIN, LOW); // Turn off LED after 4 seconds without feedback
  } else {
    digitalWrite(LED_BUILTIN, HIGH); // Turn on LED when feedback is received
  }
}