// ===========================
// Wireless Rocket Launch System v2
// Implements: MOSFET output, interrupts, state machine, non-blocking scheduler,
// ESP-NOW encryption, CRC16, sequence numbering, OTA updates, watchdog, heartbeat,
// RSSI monitoring, Adafruit SSD1306 OLED UI.
// ===========================

#include <Arduino.h>
#include <WiFi.h>
#include <esp_now.h>
#include <esp_task_wdt.h>
#include <ArduinoOTA.h>
#include <Adafruit_SSD1306.h>
#include <Wire.h>

// ======= Pin Definitions =======
// Transmitter pins
#define TX_PIN_ARM_SWITCH   14  // toggle, interrupt
#define TX_PIN_LAUNCH_BTN   15  // push, interrupt
#define TX_PIN_POWER_BTN    16  // push, interrupt
#define TX_LED_PWR          26
#define TX_LED_ARMED        12
#define TX_LED_CONN         27
#define TX_LED_PULSE        13
#define TX_OLED_SDA         SDA
#define TX_OLED_SCL         SCL

// Receiver pins
#define RX_PIN_MOSFET1      14  // MOSFET channel 1
#define RX_PIN_MOSFET2      25  // MOSFET channel 2 (redundant)
#define RX_PIN_CONTINUITY   15  // analog sense via ADC
#define RX_PIN_PWR_BTN      17  // momentary, interrupt
#define RX_LED_PWR          26
#define RX_LED_CONN         27
#define RX_OLED_SDA         SDA
#define RX_OLED_SCL         SCL

// OLED settings
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire);

// ======= ESP-NOW Encryption Key =======
uint8_t espnow_key[16] = {
  0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,
  0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F,0x10
};

// ======= Packet Structure =======
struct __attribute__((packed)) CommandMessage {
  bool arm;
  bool launch;
  uint16_t seq;
  uint16_t crc;
};

// ======= Globals =======
volatile bool armRequested = false;
volatile bool launchRequested = false;
volatile bool powerOn = false;
uint16_t txSeq = 0;
unsigned long lastSend = 0;

// Transmitter MAC (to be filled with receiver MAC)
uint8_t peerMac[6] = {0xFF,0xFF,0xFF,0xFF,0xFF,0xFF};

// ======= CRC16-CCITT =======
uint16_t crc16_ccitt(const uint8_t *data, size_t len) {
  uint16_t crc = 0xFFFF;
  while (len--) {
    crc ^= (uint16_t)*data++ << 8;
    for (uint8_t i = 0; i < 8; i++) {
      if (crc & 0x8000) crc = (crc << 1) ^ 0x1021;
      else crc <<= 1;
    }
  }
  return crc;
}

// ======= Interrupt Handlers =======
void IRAM_ATTR handleArmSwitch() {
  armRequested = (digitalRead(TX_PIN_ARM_SWITCH) == LOW);
}

void IRAM_ATTR handleLaunchBtn() {
  if (armRequested && digitalRead(TX_PIN_LAUNCH_BTN) == LOW) {
    launchRequested = true;
  }
}

void IRAM_ATTR handlePowerBtn() {
  powerOn = (digitalRead(TX_PIN_POWER_BTN) == LOW);
}

// ======= ESP-NOW Send Callback =======
void onDataSent(const uint8_t *mac_addr, esp_now_send_status_t status) {
  digitalWrite(TX_LED_CONN, status == ESP_NOW_SEND_SUCCESS);
}

// ======= Setup =======
void setup() {
  // Serial & Watchdog
  Serial.begin(115200);
  esp_task_wdt_init(5, true);
  esp_task_wdt_add(NULL);

  // Pins
  pinMode(TX_PIN_ARM_SWITCH, INPUT_PULLUP);
  pinMode(TX_PIN_LAUNCH_BTN, INPUT_PULLUP);
  pinMode(TX_PIN_POWER_BTN, INPUT_PULLUP);
  pinMode(TX_LED_PWR, OUTPUT);
  pinMode(TX_LED_ARMED, OUTPUT);
  pinMode(TX_LED_CONN, OUTPUT);
  pinMode(TX_LED_PULSE, OUTPUT);

  // Attach interrupts
  attachInterrupt(TX_PIN_ARM_SWITCH, handleArmSwitch, CHANGE);
  attachInterrupt(TX_PIN_LAUNCH_BTN, handleLaunchBtn, FALLING);
  attachInterrupt(TX_PIN_POWER_BTN, handlePowerBtn, CHANGE);

  // OLED
  Wire.begin(TX_OLED_SDA, TX_OLED_SCL);
  if (display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0,0);
    display.println("Transmitter v2 ready");
    display.display();
  }

  // Power LED
  if (!powerOn) {
    Serial.println("Power button not pressed. Halting.");
    while (!powerOn) esp_task_wdt_reset();
  }
  digitalWrite(TX_LED_PWR, HIGH);

  // Wi‑Fi & ESP-NOW
  WiFi.mode(WIFI_STA);
  esp_now_init();
  esp_now_register_send_cb(onDataSent);

  esp_now_peer_info_t peer = {};
  memcpy(peer.peer_addr, peerMac, 6);
  peer.channel = 0;
  peer.encrypt = true;
  memcpy(peer.lmk, espnow_key, 16);
  esp_now_add_peer(&peer);

  // OTA
  ArduinoOTA.begin();
}

// ======= Main Loop =======
void loop() {
  esp_task_wdt_reset();

  // OTA handling
  ArduinoOTA.handle();

  // State updates
  digitalWrite(TX_LED_ARMED, armRequested);

  // Send heartbeat/command every 100ms
  if (millis() - lastSend >= 100) {
    lastSend = millis();
    CommandMessage msg;
    msg.arm = armRequested;
    msg.launch = launchRequested;
    msg.seq = txSeq++;
    msg.crc = crc16_ccitt((uint8_t*)&msg, sizeof(msg) - sizeof(msg.crc));

    esp_now_send(peerMac, (uint8_t*)&msg, sizeof(msg));
    launchRequested = false;

    // Pulse LED on launch request
    if (msg.launch) {
      analogWrite(TX_LED_PULSE, 255);
    } else {
      analogWrite(TX_LED_PULSE, 0);
    }

    // Optionally read and display RSSI
    int8_t rssi;
    esp_now_fetch_peer(peerMac, nullptr);
    esp_now_get_peer_rssi(peerMac, &rssi);
    display.clearDisplay();
    display.setCursor(0,0);
    display.printf("ARM:%d LAUNCH:%d", msg.arm, msg.launch);
    display.setCursor(0,10);
    display.printf("RSSI:%ddBm", rssi);
    display.display();
  }
}