code.cpp

/////////////////////////
// server.cpp
/////////////////////////

#include <Arduino.h>
#include <WiFi.h>
#include <ArduinoJson.h>
#include <ESPmDNS.h>
#include <esp_sntp.h>
#include <PsychicHttp.h>
#include <freertos/queue.h>
#include <map>
#include "server.h"

// Forward declarations
extern const char *ssid;
extern const char *password;
extern const char *local_hostname;

PsychicHttpServer server;
PsychicWebSocketHandler websocketHandler;

typedef struct {
  int socket;
  char *buffer;
  size_t len;
} WebsocketMessage;

QueueHandle_t wsMessages;

// Add near the top with other global variables

// Add near other global variables
static std::map<String, MessageHandler> messageHandlers;

void listSSIDs() {
  int n = WiFi.scanNetworks();
  Serial.println("SSIDs found:");
  for (int i = 0; i < n; i++) {
    Serial.println(WiFi.SSID(i));
  }
}

bool connectToWifi() {
    Serial.print("[WiFi] Connecting to ");
    Serial.println(ssid);

    WiFi.mode(WIFI_STA);
    WiFi.begin(ssid, password);
    WiFi.setSleep(false);  // Disable WiFi power save mode
    

    int tryDelay = 500;
    int numberOfTries = 20;

    while (true) {
        switch (WiFi.status()) {
            case WL_NO_SSID_AVAIL:
                Serial.println("[WiFi] SSID not found");
                break;
            case WL_CONNECT_FAILED:
                Serial.print("[WiFi] Failed - WiFi not connected! Reason: ");
                return false;
            case WL_CONNECTION_LOST:
                Serial.println("[WiFi] Connection was lost");
                break;
            case WL_SCAN_COMPLETED:
                Serial.println("[WiFi] Scan is completed");
                break;
            case WL_DISCONNECTED:
                Serial.println("[WiFi] WiFi is disconnected");
                break;
            case WL_CONNECTED:
                Serial.println("[WiFi] WiFi is connected!");
                Serial.print("[WiFi] IP address: ");
                Serial.println(WiFi.localIP());
                return true;
            default:
                Serial.print("[WiFi] WiFi Status: ");
                Serial.println(WiFi.status());
                break;
        }
        vTaskDelay(tryDelay);

        if (numberOfTries <= 0) {
            Serial.print("[WiFi] Failed to connect to WiFi!");
            WiFi.disconnect();
            return false;
        } else {
            numberOfTries--;
        }
    }
    return false;
}

void setupWebsocketHandlers() {
    wsMessages = xQueueCreate(10, sizeof(WebsocketMessage));
    if (wsMessages == 0)
        Serial.printf("Failed to create queue= %p\n", wsMessages);

    websocketHandler.onOpen([](PsychicWebSocketClient *client) {
        Serial.printf("[socket] connection #%u connected from %s\n", client->socket(), client->remoteIP().toString());
        client->sendMessage("Hello!");
    });

    websocketHandler.onFrame([](PsychicWebSocketRequest *request, httpd_ws_frame *frame) {
        Serial.printf("[socket] #%d sent: %s\n", request->client()->socket(), (char *)frame->payload);

        WebsocketMessage wm;
        wm.socket = request->client()->socket();
        wm.len = frame->len;
        wm.buffer = (char *)malloc(frame->len);

        if (wm.buffer == NULL) {
            Serial.printf("Queue message: unable to allocate %d bytes\n", frame->len);
            return ESP_FAIL;    
        }

        memcpy(wm.buffer, frame->payload, frame->len); 

        if (xQueueSend(wsMessages, &wm, 1) != pdTRUE) {
            Serial.printf("[socket] queue full #%d\n", wm.socket);
            free(wm.buffer);
        }

        if (!uxQueueSpacesAvailable(wsMessages))
            return request->reply("Queue Full");

        return ESP_OK;
    });

    websocketHandler.onClose([](PsychicWebSocketClient *client) {
        Serial.printf("[socket] connection #%u closed from %s\n", client->socket(), client->remoteIP().toString());
    });
}

bool setupServer() {
    if (!MDNS.begin(local_hostname)) {
        Serial.println("Error starting mDNS");
        return false;
    }
    MDNS.addService("http", "tcp", 80);

    server.listen(80);
    setupWebsocketHandlers();
    server.on("/ws", &websocketHandler);
    return true;
}

void sendWebsocketMessage(const char* message) {
    websocketHandler.sendAll(message);
}

static StaticJsonDocument<128> sendJsonDoc;  // Static JSON document
static char sendJsonBuffer[128];             // Static buffer for JSON string

void sendJsonMessage(const char* type, const JsonVariant& data) {
    sendJsonDoc.clear();  // Clear previous contents
    sendJsonDoc["type"] = type;
    sendJsonDoc["data"] = data;
    
    size_t len = serializeJson(sendJsonDoc, sendJsonBuffer, sizeof(sendJsonBuffer));
    
    if (len > 0 && len < sizeof(sendJsonBuffer)) {
        websocketHandler.sendAll(sendJsonBuffer);
    } else {
        Serial.println("[JSON] Error serializing message - buffer too small");
    }
}

void registerWebsocketMessageHandler(const char* type, MessageHandler handler) {
    messageHandlers[String(type)] = handler;
}

static StaticJsonDocument<128> receiveJsonDoc;  // Static JSON document
static char receiveJsonBuffer[128];             // Static buffer for JSON string

void processWebSocketMessages() {
    WebsocketMessage message;
    while (xQueueReceive(wsMessages, &message, 0) == pdTRUE) {
        // Make sure our client is still good
        PsychicWebSocketClient *client = websocketHandler.getClient(message.socket);
        if (client == NULL) {
            Serial.printf("[socket] client #%d bad, bailing\n", message.socket);
            free(message.buffer);
            continue;
        }

        // Try to parse the JSON message
        receiveJsonDoc.clear();
        DeserializationError error = deserializeJson(receiveJsonDoc, message.buffer, message.len);
        
        if (error) {
            Serial.printf("[JSON] Deserialize failed: %s\n", error.c_str());
        } else {
            // Get the message type
            const char* type = receiveJsonDoc["type"];
            if (type) {
                // Look up the handler for this message type
                auto handler = messageHandlers.find(String(type));
                if (handler != messageHandlers.end()) {
                    // Call the handler with the data
                    handler->second(receiveJsonDoc["data"]);
                } else {
                    Serial.printf("[socket] No handler for message type: %s\n", type);
                }
            }
        }

        // Make sure to release our memory!
        free(message.buffer);
    }
}

////////////////
// server.h
////////////////

#pragma once

#include <ArduinoJson.h>
#include <functional>

// Define callback type for message handlers
typedef std::function<void(const JsonVariant&)> MessageHandler;

bool connectToWifi();
bool setupServer();
void listSSIDs();
void sendWebsocketMessage(const char* message);
void sendJsonMessage(const char* type, const JsonVariant& data);

// Add message handler registration
void registerWebsocketMessageHandler(const char* type, MessageHandler handler);
void processWebSocketMessages();

///////////////////
// platformio.ini
///////////////////

; PlatformIO Project Configuration File
;
;   Build options: build flags, source filter
;   Upload options: custom upload port, speed and extra flags
;   Library options: dependencies, extra library storages
;   Advanced options: extra scripting
;
; Please visit documentation for the other options and examples
; https://docs.platformio.org/page/projectconf.html

[env:esp32dev]
platform = espressif32
board = esp32dev
framework = arduino
monitor_speed = 115200
lib_deps =
    PsychicHttp

////////////////
// main.cpp
////////////////

#include <Arduino.h>
#include <ArduinoJson.h>
#include "server.h"

// Network configuration
const char *ssid = "";
const char *password = "";
const char *local_hostname = "esp32";

void handleCommand(const JsonVariant& data) {
  Serial.println("Command received: " + String(data.as<String>()));
}

void setup() {
  Serial.begin(115200);
  delay(10);
  if (connectToWifi()) {
    setupServer();
    registerWebsocketMessageHandler("command", handleCommand);
  }
}

void loop() {
  processWebSocketMessages();
}