#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);

// Transmitter MAC Address (replace with actual MAC)
uint8_t transmitterMacAddress[] = {0x5C, 0x01, 0x3B, 0x6B, 0x68, 0x48};

// Pin Definitions
#define MOSFET_PIN 14      // Controls the MOSFET for igniter
#define CONTINUITY_PIN 15  // Checks igniter continuity
#define POWER_LED 2        // Indicates power is on
#define CONNECTION_LED 13  // Indicates connection with transmitter
#define LAUNCH_LED 4       // Indicates launch sequence active
#define BUZZER_PIN 5       // Buzzer for pre-launch warning

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

typedef struct {
  bool continuity;
} StatusMessage;

CommandMessage command = {false, false};
StatusMessage status;
esp_now_peer_info_t peerInfo;
unsigned long lastReceivedTime = 0;
unsigned long lastSentTime = 0;
unsigned long armTime = 0;
bool isArmed = false;
bool isLaunching = false;

// Callback for data received from transmitter
void onDataRecv(const esp_now_recv_info *recvInfo, const uint8_t *incomingData, int len) {
  memcpy(&command, incomingData, sizeof(command));
  lastReceivedTime = millis();
  Serial.println("Command Received: Arm = " + String(command.arm) + ", Launch = " + String(command.launch));
}

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

  // Initialize GPIO pins
  pinMode(MOSFET_PIN, OUTPUT);
  pinMode(CONTINUITY_PIN, INPUT);
  pinMode(POWER_LED, OUTPUT);
  pinMode(CONNECTION_LED, OUTPUT);
  pinMode(LAUNCH_LED, OUTPUT);
  pinMode(BUZZER_PIN, OUTPUT);
  digitalWrite(MOSFET_PIN, LOW);
  digitalWrite(POWER_LED, HIGH); // Turn on power LED immediately as system is powered
  digitalWrite(CONNECTION_LED, LOW);
  digitalWrite(LAUNCH_LED, LOW);
  digitalWrite(BUZZER_PIN, LOW); // Ensure buzzer starts off

  // Initialize OLED
  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 ESP-NOW
  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 receive callback
  esp_now_register_recv_cb(onDataRecv);

  // Register transmitter as peer
  memcpy(peerInfo.peer_addr, transmitterMacAddress, 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() {
  // Check continuity
  bool continuity = digitalRead(CONTINUITY_PIN) == HIGH;
  status.continuity = continuity;

  // Handle arming
  if (command.arm && continuity) {
    isArmed = true;
    armTime = millis();
  } else {
    isArmed = false;
    digitalWrite(MOSFET_PIN, LOW);
  }

  // Handle launch sequence
  if (isArmed && command.launch && !isLaunching) {
    isLaunching = true;
    digitalWrite(LAUNCH_LED, HIGH); // Launch LED ON
    digitalWrite(BUZZER_PIN, HIGH); // Buzzer ON
    delay(2000);                    // 2-second buzzer
    digitalWrite(BUZZER_PIN, LOW);  // Buzzer OFF
    if (isArmed) {                  // Check if still armed
      digitalWrite(MOSFET_PIN, HIGH); // Activate igniter
      delay(1000);                    // 1-second pulse
      digitalWrite(MOSFET_PIN, LOW);  // Deactivate igniter
    }
    digitalWrite(LAUNCH_LED, LOW);  // Launch LED OFF
    isArmed = false;
    command.launch = false;
    isLaunching = false;
  }

  // Arming timeout (30 seconds)
  if (isArmed && (millis() - armTime > 30000)) {
    isArmed = false;
  }

  // Send status to transmitter every 100ms
  if (millis() - lastSentTime >= 100) {
    esp_err_t result = esp_now_send(transmitterMacAddress, (uint8_t *)&status, sizeof(status));
    if (result == ESP_OK) {
      Serial.println("Status Sent: Continuity = " + String(status.continuity));
    } else {
      Serial.println("Status Send Failed");
    }
    lastSentTime = millis();
  }

  // Update connection LED
  if (millis() - lastReceivedTime < 1000) {
    digitalWrite(CONNECTION_LED, HIGH); // Connected
  } else {
    digitalWrite(CONNECTION_LED, LOW);  // Not connected
  }

  // Update OLED display
  static unsigned long lastUpdateTime = 0;
  if (millis() - lastUpdateTime >= 100) {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("System Powered"));
    display.print(F("Connection: "));
    display.println((millis() - lastReceivedTime < 1000) ? F("Connected") : F("No Connection"));
    display.print(F("Continuity: "));
    display.println(continuity ? F("OK") : F("Fail"));
    display.print(F("Status: "));
    display.println(isArmed ? F("Armed") : F("Disarmed"));
    if (isLaunching) {
      display.println(F("Igniting..."));
    }
    display.display();
    lastUpdateTime = millis();
  }

  delay(50); // Loop runs every 50ms
}