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

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

// Receiver MAC Address (replace with actual MAC)
uint8_t receiverMacAddress[] = {0x5C, 0x01, 0x3B, 0x6C, 0x60, 0x70};

// Pin Definitions
#define ARM_SWITCH 14
#define LAUNCH_BUTTON 15
#define POWER_LED 2
#define ARMED_LED 12
#define CONNECTION_LED 13

// ESP-NOW Data Structures
typedef struct {
  bool arm;
  bool launch;
} CommandMessage;

typedef struct {
  bool continuity;
} StatusMessage;

CommandMessage command = {false, false};
StatusMessage status = {false}; // Explicit initialization
esp_now_peer_info_t peerInfo;
bool isConnected = false;
bool armSwitch = false;
bool isIgnitionGo = false;
bool isSystemGo = false;
unsigned long lastStatusTime = 0;
unsigned long launchButtonPressStart = 0;
bool launchButtonHeld = false;
const unsigned long LAUNCH_HOLD_TIME = 2000; // 2 seconds

// Callback for send status
void onDataSent(const uint8_t *mac_addr, esp_now_send_status_t status) {
  isConnected = (status == ESP_NOW_SEND_SUCCESS);
  digitalWrite(CONNECTION_LED, isConnected ? HIGH : LOW);
  Serial.println("onDataSent: isConnected = " + String(isConnected));
}

// Callback for received status
void onDataRecv(const esp_now_recv_info *recvInfo, const uint8_t *incomingData, int len) {
  if (incomingData != nullptr && len == sizeof(StatusMessage)) {
    memcpy(&status, incomingData, sizeof(status));
    lastStatusTime = millis();
    Serial.println("onDataRecv: Continuity = " + String(status.continuity));
  } else {
    Serial.println("onDataRecv: Invalid data received");
  }
}

void setup() {
  Serial.begin(115200);
  
  // Initialize pins
  pinMode(ARM_SWITCH, INPUT_PULLUP);
  pinMode(LAUNCH_BUTTON, INPUT_PULLUP);
  pinMode(POWER_LED, OUTPUT);
  pinMode(ARMED_LED, OUTPUT);
  pinMode(CONNECTION_LED, OUTPUT);
  digitalWrite(POWER_LED, HIGH); // Power LED on as system is powered
  digitalWrite(ARMED_LED, LOW);
  digitalWrite(CONNECTION_LED, LOW);

  // Initialize I2C and OLED
  Wire.begin();
  Wire.setClock(100000); // Set I2C clock to 100kHz
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    while (true);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println(F("System Powered"));
  display.display();

  // Initialize Wi-Fi
  WiFi.mode(WIFI_STA);
  if (esp_now_init() != ESP_OK) {
    Serial.println("ESP-NOW Init Failed");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("ESP-NOW Init Failed"));
    display.display();
    while (true);
  }

  // Register callbacks
  esp_now_register_send_cb(onDataSent);
  esp_now_register_recv_cb(onDataRecv);

  // Register peer
  memcpy(peerInfo.peer_addr, receiverMacAddress, 6);
  peerInfo.channel = 0;
  peerInfo.encrypt = false;
  if (esp_now_add_peer(&peerInfo) != ESP_OK) {
    Serial.println("Peer Add Failed");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("Peer Add Failed"));
    display.display();
    while (true);
  }
}

void loop() {
  static unsigned long lastUpdateTime = 0;

  Serial.println("Reading arm switch...");
  // Read arming switch
  bool armSwitch = digitalRead(ARM_SWITCH) == LOW;
  Serial.println("Arm switch: " + String(armSwitch));
  isIgnitionGo = status.continuity && (millis() - lastStatusTime < 1000);
  bool armStatus = armSwitch && isIgnitionGo;
  if(isIgnitionGo){
    command.arm = armSwitch; // Send arm intent regardless of continuity
  }
  Serial.println("Arm status: " + String(armStatus) + ", Continuity: " + String(status.continuity));
  if(armStatus && isConnected && isIgnitionGo) {
      isSystemGo = true;
    }
  // Read launch button with 2-second hold requirement
  bool launchButtonPressed = digitalRead(LAUNCH_BUTTON) == LOW;
  Serial.println("Launch button: " + String(launchButtonPressed));

  if (launchButtonPressed) {
    if (!launchButtonHeld) {
      // Button just pressed
      launchButtonPressStart = millis();
      launchButtonHeld = true;
      Serial.println("Launch button pressed, starting hold timer");
    } else if (millis() - launchButtonPressStart >= LAUNCH_HOLD_TIME) {
      // Button held for 2 seconds
      command.launch = armStatus; // Only launch if armed
      Serial.println("Launch command set: " + String(command.launch));
    }
  } else {
    // Button released or not pressed
    launchButtonHeld = false;
    command.launch = false;
    Serial.println("Launch button released or not pressed");
  }

  // Update armed LED
  digitalWrite(ARMED_LED, armStatus ? HIGH : LOW);

  // Update OLED display
  if (millis() - lastUpdateTime >= 100) {
    Serial.println("Updating OLED...");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("GO/NO-GO Status"));
    display.print(F("Arming: "));
    display.println(armStatus ? F("GO") : F("NO-GO"));
    display.print(F("Receiver Link: "));
    display.println(isConnected ? F("GO") : F("NO-GO"));
    display.print(F("Ignition: "));
    display.println(isIgnitionGo ? F("GO") : F("NO-GO"));
    display.print(F("System: "));
    display.print(isSystemGo ? F("GO") : F("NO-GO"));
    if (command.launch) {
      display.println(F("Launching..."));
    }
    display.display();
    lastUpdateTime = millis();
    Serial.println("OLED updated");
  }

  // Send command
  Serial.println("Sending command...");
  esp_err_t result = esp_now_send(receiverMacAddress, (uint8_t *)&command, sizeof(command));
  if (result == ESP_OK) {
    Serial.println("Command Sent");
  } else {
    Serial.println("Send Failed");
  }

  delay(100); // Update every 100ms
}