// ===========================
// 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.
// ===========================

// ===========================
// ===== RECEIVER CODE ======
// ===========================

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

// Globals
volatile bool powerOnRx = false;
uint16_t lastSeq = 0;
unsigned long lastRecv = 0;
bool isArmed = false;

// OLED
Adafruit_SSD1306 oled(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire);

// CRC
uint16_t crc16_ccitt_rx(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
void IRAM_ATTR handlePowerBtnRx() {
  powerOnRx = (digitalRead(RX_PIN_PWR_BTN) == LOW);
}

// Receive callback
void onDataRecvRx(const uint8_t *mac, const uint8_t *incomingData, int len) {
  struct CommandMessage {
    bool arm; bool launch; uint16_t seq; uint16_t crc;
  } msg;
  memcpy(&msg, incomingData, sizeof(msg));
  uint16_t calc = crc16_ccitt_rx((uint8_t*)&msg, sizeof(msg) - sizeof(msg.crc));
  if (calc != msg.crc || msg.seq <= lastSeq) return;  // drop bad/replay
  lastSeq = msg.seq;
  lastRecv = millis();
  isArmed = msg.arm;
  if (isArmed && msg.launch) {
    // fire mosfets
    digitalWrite(RX_PIN_MOSFET1, HIGH);
    digitalWrite(RX_PIN_MOSFET2, HIGH);
    delay(1000);
    digitalWrite(RX_PIN_MOSFET1, LOW);
    digitalWrite(RX_PIN_MOSFET2, LOW);
    isArmed = false;
  }
}

void setupRx() {
  Serial.begin(115200);
  esp_task_wdt_init(5, true);
  esp_task_wdt_add(NULL);

  // Pins
  pinMode(RX_PIN_MOSFET1, OUTPUT);
  pinMode(RX_PIN_MOSFET2, OUTPUT);
  pinMode(RX_PIN_CONTINUITY, INPUT);
  pinMode(RX_PIN_PWR_BTN, INPUT_PULLUP);
  pinMode(RX_LED_PWR, OUTPUT);
  pinMode(RX_LED_CONN, OUTPUT);
  attachInterrupt(RX_PIN_PWR_BTN, handlePowerBtnRx, CHANGE);

  // OLED
  Wire.begin(RX_OLED_SDA, RX_OLED_SCL);
  if (oled.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    oled.clearDisplay(); oled.setTextSize(1);
    oled.setTextColor(SSD1306_WHITE);
    oled.setCursor(0,0);
    oled.println("Receiver v2 ready"); oled.display();
  }

  // Wait for power
  if (!powerOnRx) {
    Serial.println("RX Power off. Halt.");
    while (!powerOnRx) esp_task_wdt_reset();
  }
  digitalWrite(RX_LED_PWR, HIGH);

  // Wi-Fi & ESP-NOW
  WiFi.mode(WIFI_STA);
  esp_now_init();
  esp_now_register_recv_cb(onDataRecvRx);
  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();
}

void loopRx() {
  esp_task_wdt_reset();
  ArduinoOTA.handle();

  // Connection LED
  digitalWrite(RX_LED_CONN, millis() - lastRecv < 1000);

  // Continuity analog
  int cont = analogRead(RX_PIN_CONTINUITY);
  if (!cont || !isArmed) {
    isArmed = false;
  }

  // Timeout
  static unsigned long armTime = 0;
  if (isArmed && armTime == 0) armTime = lastRecv;
  if (isArmed && millis() - armTime > 30000) isArmed = false;
  if (!isArmed) armTime = 0;

  // OLED status
  oled.clearDisplay();
  oled.setCursor(0,0);
  oled.printf("Conn:%d Armed:%d", digitalRead(RX_LED_CONN), isArmed);
  oled.setCursor(0,10);
  oled.printf("Cont:%d", cont);
  oled.display();

  delay(50);
}

void main() {
  setupRx();
  while (true) loopRx();
}