/*
============================================================================
PLC-GRADE TUYA ESP32 RELAY CONTROLLER
FACTORY PRODUCTION EDITION - REV 15.0
============================================================================

ARCHITECTURE:
- Dual Core FreeRTOS
- Dedicated Tuya Task
- Dedicated NVS Task
- Non-blocking WiFi Recovery
- Deterministic Relay Engine
- Industrial Safe Mode
- Queue-based Relay Commands
- Lock-protected Shared State
- Brownout Aware
- Flash Endurance Optimized
- Event-driven Runtime
- Zero Blocking Main Loop

READY FOR 24/7 DEPLOYMENT
============================================================================
*/

#include <Arduino.h>
#include <WiFi.h>
#include <TuyaWifi.h>
#include <Preferences.h>

#include <esp_task_wdt.h>
#include <esp_reset_reason.h>

#include "driver/gpio.h"

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/queue.h"
#include "freertos/semphr.h"

// ============================================================================
// SYSTEM
// ============================================================================

enum SystemState {
  STATE_BOOT,
  STATE_NORMAL,
  STATE_SAFE_MODE,
  STATE_PAIRING,
  STATE_WIFI_RECOVERY
};

enum FaultCode {
  FAULT_NONE,
  FAULT_WIFI_TIMEOUT,
  FAULT_HEAP_CRITICAL,
  FAULT_RELAY_MISMATCH,
  FAULT_NVS_CORRUPT,
  FAULT_TUYA_STALL,
  FAULT_LOOP_BLOCKED,
  FAULT_BROWNOUT
};

enum RelayCommandType {
  CMD_SET_RELAY
};

struct RelayCommand {
  RelayCommandType type;
  uint8_t id;
  bool state;
};

// ============================================================================
// CONFIG
// ============================================================================

#define FW_NAME                        "PLC-TUYA-RELAY"
#define FW_VERSION                     "15.0.0"

#define DEBUG                          true

#define NUM_RELAYS                     4
#define RELAY_ACTIVE_LOW               true

#define STATUS_LED_PIN                 4
#define PAIR_BUTTON_PIN                27

#define ENABLE_RELAY_FEEDBACK          false

#define RELAY_MIN_INTERVAL_MS          250

#define FLASH_SAVE_INTERVAL_MS         15000UL
#define FLASH_STABLE_TIME_MS           30000UL

#define WIFI_RECOVERY_COOLDOWN_MS      45000UL
#define WIFI_WATCHDOG_TIMEOUT_MS       180000UL

#define TUYA_HEALTH_TIMEOUT_MS         30000UL

#define RELAY_FEEDBACK_TIMEOUT_MS      1000UL

#define BUTTON_DEBOUNCE_MS             50

#define STABLE_UPTIME_MS               60000UL

#define MAX_BOOT_COUNT                 5

#define HEAP_CRITICAL_THRESHOLD        10000
#define MAX_ALLOC_CRITICAL             4096

#define AP_MODE_TIMEOUT_MS             300000UL

#define RELAY_RESTORE_MODE             1

// ============================================================================
// GPIO
// ============================================================================

const gpio_num_t relayGPIOs[NUM_RELAYS] = {
  GPIO_NUM_26,
  GPIO_NUM_25,
  GPIO_NUM_32,
  GPIO_NUM_33
};

const gpio_num_t relayFeedbackPins[NUM_RELAYS] = {
  GPIO_NUM_34,
  GPIO_NUM_35,
  GPIO_NUM_36,
  GPIO_NUM_39
};

const char* const NVS_KEYS[NUM_RELAYS] = {
  "r0",
  "r1",
  "r2",
  "r3"
};

// ============================================================================
// TUYA
// ============================================================================

TuyaWifi my_device;

unsigned char pid[] = "YOUR_PRODUCT_ID_HERE";
unsigned char mcu_ver[] = "15.0.0";

#define DPID_SWITCH_1 1
#define DPID_SWITCH_2 2
#define DPID_SWITCH_3 3
#define DPID_SWITCH_4 4

unsigned char dp_array[][2] = {
  {DPID_SWITCH_1, DP_TYPE_BOOL},
  {DPID_SWITCH_2, DP_TYPE_BOOL},
  {DPID_SWITCH_3, DP_TYPE_BOOL},
  {DPID_SWITCH_4, DP_TYPE_BOOL}
};

// ============================================================================
// GLOBALS
// ============================================================================

Preferences prefs;

SemaphoreHandle_t relayMutex;

QueueHandle_t relayQueue;

// ============================================================================
// STATES
// ============================================================================

bool relayState[NUM_RELAYS];
bool relayDirty[NUM_RELAYS];
bool cloudDirty[NUM_RELAYS];

uint32_t relayVersion[NUM_RELAYS];

uint32_t lastRelaySwitch[NUM_RELAYS];
uint32_t relayCommandTime[NUM_RELAYS];

// ============================================================================
// SYSTEM STATUS
// ============================================================================

volatile SystemState currentState = STATE_BOOT;
volatile FaultCode lastFault = FAULT_NONE;

bool apModeActive = false;

bool systemReady = false;

uint32_t bootCounter = 0;

// ============================================================================
// TIMERS
// ============================================================================

uint32_t lastFlashSave = 0;

uint32_t lastStateStableTime = 0;

uint32_t wifiDisconnectedSince = 0;

uint32_t lastWiFiRecovery = 0;

uint32_t lastTuyaPacketTime = 0;

uint32_t apModeStartTime = 0;

// ============================================================================
// DEBUG
// ============================================================================

#if DEBUG
#define LOG(x) Serial.println(x)
#define LOGF(...) Serial.printf(__VA_ARGS__)
#else
#define LOG(x)
#define LOGF(...)
#endif

// ============================================================================
// CRC
// ============================================================================

uint32_t calculateCRC(const bool *states) {

  uint32_t crc = 0xA5A5A5A5;

  for (uint8_t i = 0; i < NUM_RELAYS; i++) {

    crc ^= states[i]
      ? (0xFFFFFFFFU << i)
      : (0x12345678U >> i);
  }

  return crc;
}

// ============================================================================
// GPIO HAL
// ============================================================================

inline void relayWriteHardware(uint8_t id,
                               bool state) {

  gpio_set_level(
    relayGPIOs[id],
    RELAY_ACTIVE_LOW
      ? !state
      : state
  );
}

// ============================================================================
// SAFE MODE
// ============================================================================

void enterSafeMode(FaultCode reason) {

  lastFault = reason;

  currentState = STATE_SAFE_MODE;

  for (uint8_t i = 0; i < NUM_RELAYS; i++) {

    relayWriteHardware(i, false);
  }

  LOG("[SAFE MODE] Outputs disabled");
}

// ============================================================================
// STORAGE
// ============================================================================

bool loadRelayStates() {

  bool temp[NUM_RELAYS];

  for (uint8_t i = 0; i < NUM_RELAYS; i++) {

    temp[i] =
      prefs.getBool(NVS_KEYS[i], false);
  }

  uint32_t storedCRC =
    prefs.getUInt("crc", 0);

  uint32_t calcCRC =
    calculateCRC(temp);

  if (storedCRC != 0 &&
      storedCRC != calcCRC) {

    LOG("[NVS] CRC INVALID");

    enterSafeMode(FAULT_NVS_CORRUPT);

    memset(relayState, 0, sizeof(relayState));

    return false;
  }

  memcpy(relayState,
         temp,
         sizeof(temp));

  LOG("[NVS] CRC VALID");

  return true;
}

void saveRelayStates() {

  uint32_t now = millis();

  if (now - lastFlashSave <
      FLASH_SAVE_INTERVAL_MS)
    return;

  if (now - lastStateStableTime <
      FLASH_STABLE_TIME_MS)
    return;

  bool snapshot[NUM_RELAYS];
  bool dirty[NUM_RELAYS];

  if (xSemaphoreTake(relayMutex,
      pdMS_TO_TICKS(50)) != pdTRUE)
    return;

  memcpy(snapshot,
         relayState,
         sizeof(snapshot));

  memcpy(dirty,
         relayDirty,
         sizeof(dirty));

  memset(relayDirty,
         0,
         sizeof(relayDirty));

  xSemaphoreGive(relayMutex);

  bool changed = false;

  for (uint8_t i = 0; i < NUM_RELAYS; i++) {

    if (!dirty[i])
      continue;

    bool oldValue =
      prefs.getBool(
        NVS_KEYS[i],
        !snapshot[i]
      );

    if (oldValue != snapshot[i]) {

      prefs.putBool(
        NVS_KEYS[i],
        snapshot[i]
      );

      changed = true;
    }
  }

  if (changed) {

    uint32_t crc =
      calculateCRC(snapshot);

    prefs.putUInt("crc", crc);

    prefs.putUInt("version", 1500);

    LOG("[NVS] Saved");
  }

  lastFlashSave = now;
}

// ============================================================================
// RELAY ENGINE
// ============================================================================

void applyRelayState(uint8_t id,
                     bool state) {

  uint32_t now = millis();

  if (now - lastRelaySwitch[id] <
      RELAY_MIN_INTERVAL_MS)
    return;

  if (xSemaphoreTake(relayMutex,
      pdMS_TO_TICKS(50)) != pdTRUE)
    return;

  if (relayState[id] != state) {

    relayState[id] = state;

    relayDirty[id] = true;

    cloudDirty[id] = true;

    relayVersion[id]++;

    relayCommandTime[id] = now;

    lastRelaySwitch[id] = now;

    lastStateStableTime = now;

    relayWriteHardware(id, state);

    LOGF("[RELAY] CH%d -> %s\n",
         id + 1,
         state ? "ON" : "OFF");
  }

  xSemaphoreGive(relayMutex);
}

// ============================================================================
// CLOUD SYNC
// ============================================================================

void handleCloudSync() {

  static uint32_t lastSync = 0;

  if (millis() - lastSync < 500)
    return;

  lastSync = millis();

  bool snapshot[NUM_RELAYS];
  bool dirty[NUM_RELAYS];

  uint32_t versions[NUM_RELAYS];

  if (xSemaphoreTake(relayMutex,
      pdMS_TO_TICKS(20)) != pdTRUE)
    return;

  memcpy(snapshot,
         relayState,
         sizeof(snapshot));

  memcpy(dirty,
         cloudDirty,
         sizeof(dirty));

  memcpy(versions,
         relayVersion,
         sizeof(versions));

  xSemaphoreGive(relayMutex);

  for (uint8_t i = 0; i < NUM_RELAYS; i++) {

    if (!dirty[i])
      continue;

    my_device.mcu_dp_update(
      i + 1,
      snapshot[i] ? 1 : 0,
      1
    );

    if (xSemaphoreTake(relayMutex,
        pdMS_TO_TICKS(20)) == pdTRUE) {

      if (relayVersion[i] ==
          versions[i]) {

        cloudDirty[i] = false;
      }

      xSemaphoreGive(relayMutex);
    }

    break;
  }
}

// ============================================================================
// FEEDBACK
// ============================================================================

bool validateRelayFeedback(uint8_t id) {

#if ENABLE_RELAY_FEEDBACK

  bool expected = relayState[id];

  bool actual =
    gpio_get_level(
      relayFeedbackPins[id]
    );

  return expected == actual;

#else

  return true;

#endif
}

// ============================================================================
// WIFI RECOVERY FSM
// ============================================================================

void handleWiFiRecovery() {

  static uint8_t step = 0;
  static uint32_t timer = 0;

  uint32_t now = millis();

  switch (step) {

    case 0:

      WiFi.disconnect(true, true);

      WiFi.mode(WIFI_OFF);

      timer = now;

      step = 1;

      break;

    case 1:

      if (now - timer < 300)
        return;

      WiFi.mode(WIFI_STA);

      WiFi.setSleep(false);

      my_device.mcu_set_wifi_mode(
        SMART_CONFIG
      );

      lastWiFiRecovery = now;

      currentState = STATE_NORMAL;

      step = 0;

      LOG("[WIFI] Recovery complete");

      break;
  }
}

// ============================================================================
// BUTTON
// ============================================================================

void handlePairButton() {

  static bool lastBtn = HIGH;
  static bool holding = false;

  static uint32_t debounce = 0;
  static uint32_t pressStart = 0;

  bool reading =
    digitalRead(PAIR_BUTTON_PIN);

  if (reading != lastBtn)
    debounce = millis();

  if (millis() - debounce >
      BUTTON_DEBOUNCE_MS) {

    if (reading == LOW) {

      if (!holding) {

        holding = true;

        pressStart = millis();
      }

      uint32_t hold =
        millis() - pressStart;

      if (hold > 10000) {

        LOG("[SYSTEM] Factory reset");

        prefs.clear();

        prefs.end();

        ESP.restart();
      }

      else if (hold > 5000 &&
               currentState ==
               STATE_NORMAL) {

        LOG("[PAIRING] AP mode");

        apModeActive = true;

        apModeStartTime = millis();

        currentState = STATE_PAIRING;

        my_device.mcu_set_wifi_mode(
          AP_CONFIG
        );
      }

    } else {

      holding = false;
    }
  }

  lastBtn = reading;
}

// ============================================================================
// TUYA CALLBACK
// ============================================================================

unsigned char dp_process(
  unsigned char dpid,
  const unsigned char value[],
  unsigned short length
) {

  bool state =
    my_device.mcu_get_dp_download_data(
      dpid,
      value,
      length
    );

  uint8_t id = dpid - 1;

  RelayCommand cmd;

  cmd.type = CMD_SET_RELAY;
  cmd.id = id;
  cmd.state = state;

  xQueueSend(relayQueue,
             &cmd,
             0);

  lastTuyaPacketTime = millis();

  return 0;
}

// ============================================================================
// TASKS
// ============================================================================

void relayTask(void *pv) {

  RelayCommand cmd;

  for (;;) {

    if (xQueueReceive(relayQueue,
        &cmd,
        pdMS_TO_TICKS(50)) == pdTRUE) {

      switch (cmd.type) {

        case CMD_SET_RELAY:

          applyRelayState(
            cmd.id,
            cmd.state
          );

          break;
      }
    }

    esp_task_wdt_reset();
  }
}

void tuyaTask(void *pv) {

  for (;;) {

    if (currentState !=
        STATE_SAFE_MODE) {

      my_device.uart_service();

      handleCloudSync();
    }

    vTaskDelay(pdMS_TO_TICKS(1));
  }
}

void nvsTask(void *pv) {

  for (;;) {

    saveRelayStates();

    vTaskDelay(pdMS_TO_TICKS(1000));
  }
}

void systemTask(void *pv) {

  for (;;) {

    uint32_t now = millis();

    handlePairButton();

    // =====================================================
    // WIFI WATCHDOG
    // =====================================================

    if (!apModeActive &&
        WiFi.status() != WL_CONNECTED) {

      if (wifiDisconnectedSince == 0)
        wifiDisconnectedSince = now;

      if (now - wifiDisconnectedSince >
          WIFI_WATCHDOG_TIMEOUT_MS) {

        currentState =
          STATE_WIFI_RECOVERY;
      }

    } else {

      wifiDisconnectedSince = 0;
    }

    // =====================================================
    // WIFI RECOVERY
    // =====================================================

    if (currentState ==
        STATE_WIFI_RECOVERY) {

      handleWiFiRecovery();
    }

    // =====================================================
    // PAIRING TIMEOUT
    // =====================================================

    if (apModeActive &&
        now - apModeStartTime >
        AP_MODE_TIMEOUT_MS) {

      apModeActive = false;

      WiFi.mode(WIFI_STA);

      currentState = STATE_NORMAL;
    }

    // =====================================================
    // FEEDBACK CHECK
    // =====================================================

    if (systemReady) {

      for (uint8_t i = 0;
           i < NUM_RELAYS;
           i++) {

        if (now - relayCommandTime[i] <
            RELAY_FEEDBACK_TIMEOUT_MS)
          continue;

        if (!validateRelayFeedback(i)) {

          enterSafeMode(
            FAULT_RELAY_MISMATCH
          );
        }
      }
    }

    // =====================================================
    // TUYA HEALTH
    // =====================================================

    if (now - lastTuyaPacketTime >
        TUYA_HEALTH_TIMEOUT_MS) {

      LOG("[WATCHDOG] Tuya timeout");

      enterSafeMode(
        FAULT_TUYA_STALL
      );
    }

    // =====================================================
    // HEAP
    // =====================================================

    int freeHeap =
      ESP.getFreeHeap();

    int maxAlloc =
      ESP.getMaxAllocHeap();

    if (freeHeap <
        HEAP_CRITICAL_THRESHOLD ||
        maxAlloc <
        MAX_ALLOC_CRITICAL) {

      enterSafeMode(
        FAULT_HEAP_CRITICAL
      );
    }

    updateStatusLED();

    esp_task_wdt_reset();

    vTaskDelay(pdMS_TO_TICKS(20));
  }
}

// ============================================================================
// LED
// ============================================================================

void updateStatusLED() {

  static uint32_t lastBlink = 0;
  static bool led = false;

  uint32_t now = millis();

  if (currentState ==
      STATE_SAFE_MODE) {

    digitalWrite(
      STATUS_LED_PIN,
      (now % 400 < 200)
    );

    return;
  }

  uint8_t wifiState =
    my_device.mcu_get_wifi_work_state();

  if (wifiState ==
      WIFI_CONN_CLOUD) {

    digitalWrite(
      STATUS_LED_PIN,
      HIGH
    );

    return;
  }

  uint16_t interval =
    apModeActive
      ? 80
      : 600;

  if (now - lastBlink >= interval) {

    lastBlink = now;

    led = !led;

    digitalWrite(
      STATUS_LED_PIN,
      led
    );
  }
}

// ============================================================================
// SETUP
// ============================================================================

void setup() {

  Serial.begin(115200);

  LOG("\n================================");
  LOG(FW_NAME " v" FW_VERSION);
  LOG("PLC PRODUCTION RTOS");
  LOG("================================");

  relayMutex =
    xSemaphoreCreateMutex();

  relayQueue =
    xQueueCreate(
      16,
      sizeof(RelayCommand)
    );

  if (!relayMutex ||
      !relayQueue) {

    ESP.restart();
  }

  esp_task_wdt_config_t wdt_config = {
    .timeout_ms = 10000,
    .idle_core_mask =
      (1 << portNUM_PROCESSORS) - 1,
    .trigger_panic = true
  };

  esp_task_wdt_init(&wdt_config);

  // =====================================================
  // GPIO
  // =====================================================

  for (uint8_t i = 0;
       i < NUM_RELAYS;
       i++) {

    gpio_reset_pin(relayGPIOs[i]);

    gpio_set_direction(
      relayGPIOs[i],
      GPIO_MODE_OUTPUT
    );

    relayWriteHardware(i, false);

#if ENABLE_RELAY_FEEDBACK

    gpio_set_direction(
      relayFeedbackPins[i],
      GPIO_MODE_INPUT
    );

#endif
  }

  pinMode(STATUS_LED_PIN, OUTPUT);

  pinMode(PAIR_BUTTON_PIN,
          INPUT_PULLUP);

  // =====================================================
  // NVS
  // =====================================================

  prefs.begin("relay", false);

  esp_reset_reason_t reason =
    esp_reset_reason();

  if (reason ==
      ESP_RST_BROWNOUT) {

    enterSafeMode(
      FAULT_BROWNOUT
    );
  }

  bool abnormal =
    reason == ESP_RST_PANIC ||
    reason == ESP_RST_TASK_WDT ||
    reason == ESP_RST_WDT;

  if (abnormal) {

    bootCounter =
      prefs.getUInt(
        "bootCount",
        0
      ) + 1;

  } else {

    bootCounter = 0;
  }

  prefs.putUInt(
    "bootCount",
    bootCounter
  );

  if (bootCounter >
      MAX_BOOT_COUNT) {

    enterSafeMode(
      FAULT_LOOP_BLOCKED
    );
  }

  loadRelayStates();

  if (RELAY_RESTORE_MODE == 1) {

    memset(relayState,
           0,
           sizeof(relayState));

    memset(relayDirty,
           1,
           sizeof(relayDirty));
  }

  for (uint8_t i = 0;
       i < NUM_RELAYS;
       i++) {

    relayWriteHardware(
      i,
      relayState[i]
    );

    cloudDirty[i] = true;

    relayVersion[i] = 0;

    relayCommandTime[i] =
      millis();
  }

  // =====================================================
  // TUYA
  // =====================================================

  if (currentState !=
      STATE_SAFE_MODE) {

    WiFi.mode(WIFI_STA);

    WiFi.setSleep(false);

    my_device.init(pid,
                   mcu_ver);

    my_device.set_dp_cmd_total(
      dp_array,
      sizeof(dp_array) /
      sizeof(dp_array[0])
    );

    my_device.dp_process_func_register(
      dp_process
    );
  }

  lastTuyaPacketTime = millis();

  // =====================================================
  // TASKS
  // =====================================================

  xTaskCreatePinnedToCore(
    relayTask,
    "relayTask",
    4096,
    NULL,
    3,
    NULL,
    1
  );

  xTaskCreatePinnedToCore(
    tuyaTask,
    "tuyaTask",
    4096,
    NULL,
    2,
    NULL,
    0
  );

  xTaskCreatePinnedToCore(
    nvsTask,
    "nvsTask",
    4096,
    NULL,
    1,
    NULL,
    1
  );

  xTaskCreatePinnedToCore(
    systemTask,
    "systemTask",
    4096,
    NULL,
    2,
    NULL,
    1
  );

  systemReady = true;

  currentState =
    currentState ==
    STATE_SAFE_MODE
      ? STATE_SAFE_MODE
      : STATE_NORMAL;

  LOG("[SYSTEM] RTOS Ready");
}

// ============================================================================
// LOOP
// ============================================================================

void loop() {

  taskYIELD();
}