#include <esp_now.h>
#include <WiFi.h>
#include <ESP32Servo.h>
#include <Adafruit_SSD1306.h>

struct Message {
  int x;
  int y;
  bool isAuto;
  bool melody;
};
Message incomingData;

Servo steerSrv, leftSrv, rightSrv;
Adafruit_SSD1306 display(128, 64, &Wire, -1);

unsigned long lastReceiveTime = 0;
const int PIN_YELLOW = 13;
const int PIN_RGB_R = 10;
const int PIN_RGB_G = 11;

void OnDataRecv(const uint8_t * mac, const uint8_t *data, int len) {
  memcpy(&incomingData, data, sizeof(incomingData));
  lastReceiveTime = millis(); // Reset the timer every time we get a signal
}

void setup() {
  Serial.begin(115200);
  WiFi.mode(WIFI_STA);

  pinMode(PIN_YELLOW, OUTPUT);
  pinMode(PIN_RGB_R, OUTPUT);
  pinMode(PIN_RGB_G, OUTPUT);
  
  digitalWrite(PIN_YELLOW, HIGH); // Force Yellow LED ON to show board has power

  steerSrv.attach(4);
  leftSrv.attach(5);
  rightSrv.attach(6);

  Wire.begin(1, 2); 
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
  }
  
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0,10);
  display.println("BOT POWERED ON");
  display.display();

  esp_now_init();
  esp_now_register_recv_cb(esp_now_recv_cb_t(OnDataRecv));
}

void loop() {
  unsigned long now = millis();

  // SIGNAL STRENGTH LOGIC
  if (now - lastReceiveTime < 200) {
    // Strong Signal (Data received in last 0.2s)
    digitalWrite(PIN_RGB_G, HIGH);
    digitalWrite(PIN_RGB_R, LOW);
  } else if (now - lastReceiveTime < 1000) {
    // Weak/No Signal (Orange-ish: both Red and Green on)
    digitalWrite(PIN_RGB_G, HIGH);
    digitalWrite(PIN_RGB_R, HIGH);
  } else {
    // Lost Connection (Red)
    digitalWrite(PIN_RGB_G, LOW);
    digitalWrite(PIN_RGB_R, HIGH);
    
    // Safety: Stop motors if signal is lost
    incomingData.x = 2048; 
    incomingData.y = 2048;
  }

  // DRIVE LOGIC
  int throttle = map(incomingData.y, 0, 4095, -90, 90);
  int turn = map(incomingData.x, 0, 4095, -40, 40);

  steerSrv.write(90 + turn);
  leftSrv.write(90 + throttle + turn);
  rightSrv.write(90 + throttle - turn);

  // Update Display
  display.clearDisplay();
  display.setCursor(0,0);
  display.print("RSSI: "); 
  display.println(now - lastReceiveTime < 500 ? "STRONG" : "LOST");
  display.print("MODE: "); display.println(incomingData.isAuto ? "AUTO" : "MANUAL");
  display.display();

  delay(50);
}