#include <WiFi.h>
#include <WebSocketsClient.h>
#include <ArduinoJson.h>
#include <ESP32Servo.h>
#include <Preferences.h>
#include <vector>

// Define the WebSocket server address
const char* websocket_server_host = "wss-global.undernouzen.shop";
const uint16_t websocket_server_port = 443; // Standard port for WSS
const char* websocket_path = "/ESP32UND?";

// Object to store preferences in flash memory
Preferences preferences;

// Objects for communication and control
WebSocketsClient webSocketClient;
Servo myServo;
String current_mac_address;
bool ws_connected = false;

// Dynamic pin variables
// Debug property
int debug_mode = 0; // 0 = disabled, 1 = enabled
const int builtInLed = 2; // Built-in LED of the ESP32 board

// Potentiometer monitoring variables
const int potThreshold = 50; // Minimum change to trigger an event

// Touch sensor monitoring variables
const int touchThreshold = 40; // Minimum value to detect a touch

// Structure to hold sensor information for monitoring
struct MonitoredSensor {
  int pin;
  String type;
  int lastValue;
  int lastReading; // For button debouncing state
  unsigned long lastDebounceTime = 0; // For button debouncing
};

// Structure to hold LED state information
struct LedState {
  int pin;
  int state; // HIGH or LOW
};

// Vector to hold the list of active sensors
std::vector<MonitoredSensor> activeSensors;
// Vector to hold the list of configured LEDs
std::vector<LedState> activeLeds;

int servoPin = -1;
bool is_servo_attached = false;

// Function prototypes
void startWiFi(const char* ssid, const char* password);
void startWiFiAP();
void saveWiFiCredentials(const char* ssid, const char* password);
void loadWiFiCredentials(String &ssid, String &password);
void saveDebugMode(int mode);
int loadDebugMode();
String getMacAddressString();
String getMacAddressNoColons();
void webSocketEvent(WStype_t type, uint8_t * payload, size_t length);
void handleCommand(String payload, bool isSerial);
void processSerial();
void connectToWebSocket();
void blinkBuiltInLed(int duration);
void sendSensorValueToWebSocket(int pin, const char* type, int value);
void monitorActiveSensors();
void enableSensorMonitor(int pin, const char* type);
void disableSensorMonitor(int pin, const char* type);
void saveMonitorsState();
void saveLedsState();
void loadAndApplyConfigs();

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

  Serial.println("{'action':'starting'}");
  // Load debug mode before any conditional serial printing
  debug_mode = loadDebugMode();
  if (debug_mode == 1) {
    Serial.println("Debug mode: ENABLED");
  } else {
    Serial.println("Debug mode: DISABLED");
  }

  // Configure built-in LED
  pinMode(builtInLed, OUTPUT);
  digitalWrite(builtInLed, LOW);

  // Initialize WiFi mode
  WiFi.mode(WIFI_STA); 
  
  // Try to load credentials and connect
  String saved_ssid, saved_password;
  loadWiFiCredentials(saved_ssid, saved_password);

  if (saved_ssid.length() > 0 && saved_password.length() > 0) {
    if (debug_mode == 1) Serial.println("Saved WiFi credentials found. Attempting to connect...");
    startWiFi(saved_ssid.c_str(), saved_password.c_str());
  } else {
    Serial.println("No WiFi credentials saved. Starting in Access Point mode.");
    startWiFiAP();
  }

  // Load and apply saved LED and Monitor configurations
  loadAndApplyConfigs();
}

void loop() {
  // Handle WebSocket events
  if (WiFi.status() == WL_CONNECTED) {
    if (current_mac_address.length() > 0 && !ws_connected) {
        connectToWebSocket();
    }
    webSocketClient.loop();
  } else {
    ws_connected = false;
  }

  // Handle commands via Serial
  processSerial();
  
  // Monitor active sensors
  monitorActiveSensors();
  
  // A small delay to avoid overwhelming the CPU
  delay(10);
}

void startWiFi(const char* ssid, const char* password) {
  if (debug_mode == 1) Serial.println("Attempting to connect to WiFi...");
  WiFi.begin(ssid, password);

  int timeout_counter = 0;
  while (WiFi.status() != WL_CONNECTED && timeout_counter < 30) {
    delay(500);
    if (debug_mode == 1) Serial.print(".");
    timeout_counter++;
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("\nConnected to WiFi network!");
    Serial.print("IP Address: ");
    Serial.println(WiFi.localIP());

    current_mac_address = getMacAddressNoColons(); 
    Serial.print("ESP32 MAC Address (without ':'): ");
    Serial.println(current_mac_address);
  } else {
    Serial.println("\nConnection failed. Starting in Access Point mode.");
    startWiFiAP();
  }
}

void startWiFiAP() {
  Serial.println("Starting Access Point...");
  WiFi.mode(WIFI_AP);
  WiFi.softAP("UnderDeck", "12345678"); 
  Serial.print("Access Point (AP) started with SSID: ");
  Serial.println("UnderDeck");
  Serial.print("AP IP Address: ");
  Serial.println(WiFi.softAPIP());
  
  current_mac_address = getMacAddressNoColons(); 
  if (debug_mode == 1) {
    Serial.print("ESP32 MAC Address (without ':'): ");
    Serial.println(current_mac_address);
  }
}

void saveWiFiCredentials(const char* ssid, const char* password) {
  preferences.begin("wifi_creds", false);
  preferences.putString("ssid", ssid);
  preferences.putString("password", password);
  preferences.end();
  if (debug_mode == 1) Serial.println("WiFi credentials saved to memory.");
}

void loadWiFiCredentials(String &ssid, String &password) {
  preferences.begin("wifi_creds", true);
  ssid = preferences.getString("ssid", "");
  password = preferences.getString("password", "");
  preferences.end();
}

void saveDebugMode(int mode) {
  preferences.begin("system", false);
  preferences.putInt("debug_mode", mode);
  preferences.end();
  if (debug_mode == 1) Serial.println("Debug mode saved.");
}

int loadDebugMode() {
  preferences.begin("system", true);
  int mode = preferences.getInt("debug_mode", 0);
  preferences.end();
  return mode;
}

String getMacAddressString() {
  return WiFi.macAddress();
}

String getMacAddressNoColons() {
  String mac = WiFi.macAddress();
  mac.replace(":", "");
  return mac;
}

void connectToWebSocket() {
  if (ws_connected) {
    webSocketClient.disconnect();
    ws_connected = false;
  }

  if (WiFi.status() == WL_CONNECTED) {
    String url = websocket_path;
    url += current_mac_address;
    if (debug_mode == 1){
      Serial.print("Attempting to connect to wss://");
      Serial.print(websocket_server_host);
      Serial.println(url);
    }

    webSocketClient.beginSSL(websocket_server_host, websocket_server_port, url.c_str());
    webSocketClient.onEvent(webSocketEvent);
    webSocketClient.setReconnectInterval(5000);
    ws_connected = true;
  } else {
    if (debug_mode == 1) Serial.println("Not connected to WiFi, cannot connect to WebSocket.");
  }
}

void processSerial() {
  if (Serial.available()) {
    String command = Serial.readStringUntil('\n');
    command.trim();
    if (command.startsWith("{") && command.endsWith("}")) {
      if (debug_mode == 1) Serial.print("JSON command received via Serial: ");
      if (debug_mode == 1) Serial.println(command);
      blinkBuiltInLed(50);
      handleCommand(command, true);
    }
  }
}

void webSocketEvent(WStype_t type, uint8_t * payload, size_t length) {
  switch (type) {
    case WStype_CONNECTED:
      Serial.println("[WS] Connected to server!");
      ws_connected = true;
      break;
    case WStype_DISCONNECTED:
      Serial.println("[WS] Disconnected from server!");
      ws_connected = false;
      break;
    case WStype_TEXT: {
      String message = String((char*)payload);
      if (debug_mode == 1) Serial.print("[WS] Message received: ");
      if (debug_mode == 1) Serial.println(message);
      blinkBuiltInLed(50);
      handleCommand(message, false);
      break;
    }
    case WStype_BIN:
    case WStype_ERROR:
    case WStype_PONG:
      break;
  }
}

void handleCommand(String payload, bool isSerial) {
  StaticJsonDocument<512> doc; 
  DeserializationError error = deserializeJson(doc, payload);

  if (error) {
    String response_error = "{\"status\":\"error\", \"message\":\"Invalid JSON\"}";
    if (debug_mode == 1) Serial.println(response_error);
    if (!isSerial) {
      webSocketClient.sendTXT(response_error);
    }
    return;
  }
  
  const char* type = doc["type"];
  if (type != nullptr && strcmp(type, "esp32") == 0) {
      const char* code = doc["code"];
      if (code != nullptr && strcmp(code, "off-line") == 0) {
          if (debug_mode == 1) Serial.println("Received 'off-line' message from server. No action needed.");
          return; 
      }
  }

  const char* action = doc["action"];
  String response;
  StaticJsonDocument<128> responseDoc;
  responseDoc["status"] = "success";
  responseDoc["action"] = action;

  if (action == nullptr) {
    responseDoc["status"] = "error";
    responseDoc["message"] = "Action not specified or invalid JSON";
  } else if (strcmp(action, "set_debug") == 0) {
    int new_mode = doc["mode"];
    if (new_mode == 0 || new_mode == 1) {
      debug_mode = new_mode;
      saveDebugMode(new_mode);
      responseDoc["message"] = "Debug mode changed to " + String(new_mode);
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Invalid debug mode. Use 0 or 1.";
    }
  } else if (strcmp(action, "set_wifi") == 0) {
    const char* ssid = doc["ssid"];
    const char* password = doc["password"];
    if (ssid && password) {
      saveWiFiCredentials(ssid, password);
      responseDoc["message"] = "WiFi credentials saved. Restart to reconnect.";
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Missing SSID/Password parameters.";
    }
  } else if (strcmp(action, "reboot_wifi") == 0) {
    if (debug_mode == 1) Serial.println("WiFi reboot command received.");
    String saved_ssid, saved_password;
    loadWiFiCredentials(saved_ssid, saved_password);
    if (saved_ssid.length() > 0 && saved_password.length() > 0) {
        WiFi.disconnect(true);
        startWiFi(saved_ssid.c_str(), saved_password.c_str());
    } else {
        WiFi.disconnect(true);
        startWiFiAP();
    }
    responseDoc["message"] = "Attempting to restart WiFi connection.";
  } else if (strcmp(action, "get_mac") == 0) {
    responseDoc["mac_address"] = current_mac_address;
  } else if (strcmp(action, "get_status") == 0) {
    bool isConnected = WiFi.status() == WL_CONNECTED;
    responseDoc["connected"] = isConnected;
    responseDoc["message"] = isConnected ? "Connected to WiFi" : "Not connected to WiFi";
  } else if (strcmp(action, "get_device_info") == 0) {
    responseDoc["hostname"] = WiFi.getHostname();
    responseDoc["mac_address"] = current_mac_address;
  } else if (strcmp(action, "enable_monitor") == 0) {
    int pin = doc["pin"];
    const char* sensorType = doc["type"];
    if (pin >= 0 && sensorType != nullptr) {
      enableSensorMonitor(pin, sensorType);
      responseDoc["message"] = "Enabled monitoring for " + String(sensorType) + " on pin " + String(pin);
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Invalid parameters for enabling monitor.";
    }
  } else if (strcmp(action, "disable_monitor") == 0) {
    int pin = doc["pin"];
    const char* sensorType = doc["type"];
    if (pin >= 0 && sensorType != nullptr) {
      disableSensorMonitor(pin, sensorType);
      responseDoc["message"] = "Disabled monitoring for " + String(sensorType) + " on pin " + String(pin);
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Invalid parameters for disabling monitor.";
    }
  } else if (strcmp(action, "read_analog") == 0) {
    const int pin = doc["pin"];
    if (pin >= 0) {
      int sensorValue = analogRead(pin);
      sendSensorValueToWebSocket(pin, "analog", sensorValue);
      responseDoc["message"] = "Analog value read at pin " + String(pin);
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Invalid pin specified for analog sensor.";
    }
  } else if (strcmp(action, "led") == 0) {
    const int pin = doc["pin"];
    const char* state = doc["state"];
    int ledState = -1;

    if (strcmp(state, "on") == 0) ledState = HIGH;
    else if (strcmp(state, "off") == 0) ledState = LOW;

    if (pin >= 0 && ledState != -1) {
        pinMode(pin, OUTPUT);
        digitalWrite(pin, ledState);

        // Update or add the LED to our active list
        bool found = false;
        for (auto& led : activeLeds) {
            if (led.pin == pin) {
                led.state = ledState;
                found = true;
                break;
            }
        }
        if (!found) {
            activeLeds.push_back({pin, ledState});
        }
        saveLedsState(); // Save the new state to flash
        responseDoc["message"] = "LED at pin " + String(pin) + " set to " + state;
    } 
  } else if (strcmp(action, "servo") == 0) {
    const int pin = doc["pin"];
    int position = doc["position"];
    if (pin != servoPin || !is_servo_attached) {
      if (is_servo_attached) {
        myServo.detach();
      }
      myServo.attach(pin);
      servoPin = pin;
      is_servo_attached = true;
    }
    myServo.write(position);
    responseDoc["message"] = "Servo moved to " + String(position) + " degrees at pin " + String(servoPin);
  } else if (strcmp(action, "remove_led") == 0) {
    const int pin = doc["pin"];
    if (pin >= 0) {
      bool removed = false;
      for (auto it = activeLeds.begin(); it != activeLeds.end(); ++it) {
        if (it->pin == pin) {
          digitalWrite(pin, LOW); // Turn off the LED
          activeLeds.erase(it);
          removed = true;
          break;
        }
      }
      if (removed) {
        saveLedsState(); // Save the updated list to flash
        responseDoc["message"] = "LED configuration for pin " + String(pin) + " removed.";
      } else {
        responseDoc["status"] = "error";
        responseDoc["message"] = "No LED configuration found for pin " + String(pin) + ".";
      }
    } else {
      responseDoc["status"] = "error";
      responseDoc["message"] = "Invalid pin specified for remove_led.";
    }
  } else if (strcmp(action, "get_configs") == 0) {
    preferences.begin("device_state", true);
    String monitorsJson = preferences.getString("monitors", "[]");
    String ledsJson = preferences.getString("leds", "[]");
    preferences.end();

    StaticJsonDocument<512> dataDoc;
    deserializeJson(dataDoc["monitors"], monitorsJson);
    deserializeJson(dataDoc["leds"], ledsJson);
    
    // For this command, the response is the dataDoc itself
    serializeJson(dataDoc, response); 
  } else if (strcmp(action, "clear_configs") == 0) {
    // Turn off all active LEDs
    for (const auto& led : activeLeds) {
      digitalWrite(led.pin, LOW);
    }
    // Clear runtime vectors
    activeLeds.clear();
    activeSensors.clear();

    // Clear stored preferences
    preferences.begin("device_state", false);
    preferences.remove("leds");
    preferences.remove("monitors");
    preferences.end();

    if (debug_mode == 1) {
      Serial.println("All LED and monitor configurations have been cleared.");
    }
    responseDoc["message"] = "All LED and monitor configurations cleared.";
  } else {
    responseDoc["status"] = "error";
    responseDoc["message"] = "Unknown command";
  }

  // Only serialize responseDoc if the response string is still empty
  if (response.length() == 0) {
    serializeJson(responseDoc, response);
  }
  if (!isSerial) {
    webSocketClient.sendTXT(response);
    if (debug_mode == 1) Serial.println(response);
  } else {
    Serial.println(response);
  }
}

void blinkBuiltInLed(int duration) {
  digitalWrite(builtInLed, HIGH);
  delay(duration);
  digitalWrite(builtInLed, LOW);
}

void sendSensorValueToWebSocket(int pin, const char* type, int value) {
  if (ws_connected) {
    StaticJsonDocument<128> eventDoc;
    eventDoc["action"] = "sensor_update";
    eventDoc["type"] = type;
    eventDoc["pin"] = pin;
    eventDoc["value"] = value;
    eventDoc["timestamp"] = millis();
    String json_string;
    serializeJson(eventDoc, json_string);
    webSocketClient.sendTXT(json_string);
    Serial.println(json_string);
    blinkBuiltInLed(50);
    if (debug_mode == 1) Serial.println(json_string);
  } else {
    if (debug_mode == 1) {
      Serial.print("WebSocket not connected. Cannot send ");
      Serial.print(type);
      Serial.println(" sensor value.");
    }
  }
}

void monitorActiveSensors() {
  const unsigned long debounceDelay = 50;

  for (size_t i = 0; i < activeSensors.size(); ++i) {
    MonitoredSensor& sensor = activeSensors[i];
    int currentValue = 0;
    bool valueChanged = false;

    if (sensor.type == "button") {
      int currentReading = digitalRead(sensor.pin);

      // If the reading has changed, reset the debounce timer
      if (currentReading != sensor.lastReading) {
        sensor.lastDebounceTime = millis();
      }

      // If the state has been stable for long enough
      if ((millis() - sensor.lastDebounceTime) > debounceDelay) {
        // And if the stable state is different from the last reported value
        if (currentReading != sensor.lastValue) {
          currentValue = currentReading;
          valueChanged = true;
        }
      }
      sensor.lastReading = currentReading;
    } else if (sensor.type == "potentiometer") {
      currentValue = analogRead(sensor.pin);
      if (abs(currentValue - sensor.lastValue) > potThreshold) {
        valueChanged = true;
      }
    } else if (sensor.type == "touch") {
      currentValue = touchRead(sensor.pin);
      if (abs(currentValue - sensor.lastValue) > touchThreshold) {
        valueChanged = true;
      }
    }
    
    if (valueChanged) {
      sensor.lastValue = currentValue;
      sendSensorValueToWebSocket(sensor.pin, sensor.type.c_str(), currentValue);
    }
  }
}

void enableSensorMonitor(int pin, const char* type) {
  // Remove any existing monitor on the same pin to allow replacement.
  activeSensors.erase(
    std::remove_if(activeSensors.begin(), activeSensors.end(), [pin](const MonitoredSensor& sensor) {
      return sensor.pin == pin;
    }),
    activeSensors.end()
  );

  // Add sensor to the list and configure pin
  MonitoredSensor newSensor;
  newSensor.pin = pin;
  newSensor.type = type;
  
  if (newSensor.type == "button") {
    pinMode(pin, INPUT_PULLUP);
    newSensor.lastValue = digitalRead(pin);
    newSensor.lastReading = newSensor.lastValue;
  } else if (newSensor.type == "potentiometer") {
    newSensor.lastValue = analogRead(pin);
  } else if (newSensor.type == "touch") {
    newSensor.lastValue = touchRead(pin);
  }

  activeSensors.push_back(newSensor);
  saveMonitorsState(); // Save the updated list to flash
  if (debug_mode == 1) {
    Serial.print("Enabled monitoring for ");
    Serial.print(type);
    Serial.print(" on pin ");
    Serial.println(pin);
  }
}

void disableSensorMonitor(int pin, const char* type) {
  for (size_t i = 0; i < activeSensors.size(); ++i) {
    if (activeSensors[i].pin == pin && activeSensors[i].type == type) {
      activeSensors.erase(activeSensors.begin() + i);
      saveMonitorsState(); // Save the updated list to flash
      if (debug_mode == 1) {
        Serial.print("Disabled monitoring for ");
        Serial.print(type);
        Serial.print(" on pin ");
        Serial.println(pin);
      }
      return;
    }
  }
}

void saveMonitorsState() {
  StaticJsonDocument<512> doc;
  JsonArray monitorsArray = doc.to<JsonArray>();

  for (const auto& sensor : activeSensors) {
    JsonObject monitorObj = monitorsArray.createNestedObject();
    monitorObj["pin"] = sensor.pin;
    monitorObj["type"] = sensor.type;
  }

  String jsonString;
  serializeJson(doc, jsonString);

  preferences.begin("device_state", false);
  preferences.putString("monitors", jsonString);
  preferences.end();

  if (debug_mode == 1) {
    Serial.println("Saved monitors state to flash: " + jsonString);
  }
}

void saveLedsState() {
  StaticJsonDocument<256> doc;
  JsonArray ledsArray = doc.to<JsonArray>();

  for (const auto& led : activeLeds) {
    JsonObject ledObj = ledsArray.createNestedObject();
    ledObj["pin"] = led.pin;
    ledObj["state"] = (led.state == HIGH) ? "on" : "off";
  }

  String jsonString;
  serializeJson(doc, jsonString);

  preferences.begin("device_state", false);
  preferences.putString("leds", jsonString);
  preferences.end();

  if (debug_mode == 1) {
    Serial.println("Saved LEDs state to flash: " + jsonString);
  }
}

void loadAndApplyConfigs() {
  preferences.begin("device_state", true);

  // Load and apply monitors
  String monitorsJson = preferences.getString("monitors", "[]");
  StaticJsonDocument<512> monitorsDoc;
  if (deserializeJson(monitorsDoc, monitorsJson).code() == DeserializationError::Ok) {
    JsonArray monitorsArray = monitorsDoc.as<JsonArray>();
    for (JsonObject monitor : monitorsArray) {
      enableSensorMonitor(monitor["pin"], monitor["type"]);
    }
  }

  // Load and apply LED states
  String ledsJson = preferences.getString("leds", "[]");
  StaticJsonDocument<256> ledsDoc;
  if (deserializeJson(ledsDoc, ledsJson).code() == DeserializationError::Ok) {
    JsonArray ledsArray = ledsDoc.as<JsonArray>();
    for (JsonObject led : ledsArray) {
      int pin = led["pin"];
      const char* stateStr = led["state"];
      int state = (strcmp(stateStr, "on") == 0) ? HIGH : LOW;
      pinMode(pin, OUTPUT);
      digitalWrite(pin, state);
      activeLeds.push_back({pin, state});
    }
  }
  preferences.end();
}