#include "secrets.h"
#include <WiFi.h>
#include <ArduinoMqttClient.h>
#include <ArduinoJson.h>
#include <LoRa.h>
#include <SPI.h>
#include <vector>
#include <algorithm>

// --- Pin Definitions ---
#define LORA_CS 5
#define LORA_RST 14
#define LORA_IRQ 26

// --- WiFi & MQTT Configuration ---
const char* ssid = SECRET_WIFI_SSID;
const char* password = SECRET_WIFI_PASS;
const char* mqtt_server = "192.168.20.30";
const char* mqtt_user = SECRET_MQTT_USER;
const char* password = SECRET_WIFI_PASS;
const char* mqtt_topic_subscribe = "homeassistant/ESP32_RA/control";
const char* mqtt_topic_publish = "homeassistant/ESP32_RA/state";

WiFiClient wifiClient;
MqttClient mqttClient(wifiClient);

// --- LoRa Configuration ---
#define LORA_FREQUENCY 434E6

// --- Resend Queue Logic ---
struct Message {
    String message;
    String message_id;
    String targetDevice;
    unsigned long lastSentTime;
    int retryCount;
};

std::vector<Message> resendQueue;

const int resend_times = 80;
const unsigned long timerInterval = 1000;
const int effectiveIntervalSeconds = 5;
unsigned long lastTimerCheck = 0;
int sendCounter = 0;

// --- Function Declarations ---
void connectToWiFi();
void mqttReconnect();
void onMqttMessage(int messageSize);
void sendMessage(String jsonMessageToSendViaLoRa, String message_id, String targetDevice);
void checkResendQueue();
void handleReceivedLoRaMessage(String receivedMessage);
// Modified to accept JsonDocument directly and a 'retain' flag
bool publishRobustly(const char* topic, const JsonDocument& doc, bool retain, int max_attempts);


// --- Setup ---
void setup() {
    Serial.begin(115200);
    while (!Serial);
    Serial.println("\n\n--- ESP32 RA Transmitter/Gateway ---");

    // Initialize LoRa module
    Serial.println("Initializing LoRa...");
    SPI.begin();
    LoRa.setPins(LORA_CS, LORA_RST, LORA_IRQ);
    if (!LoRa.begin(LORA_FREQUENCY)) {
        Serial.println("Starting LoRa failed! Check connections/frequency.");
        while (1);
    }
    LoRa.setSpreadingFactor(9);
    LoRa.setSignalBandwidth(125E3);
    LoRa.setCodingRate4(5);
    LoRa.setPreambleLength(8);
    LoRa.setSyncWord(0x34);
    Serial.println("LoRa initialized successfully!");

    // Connect to WiFi
    connectToWiFi();

    // Configure MQTT client and callback
    mqttClient.setUsernamePassword(mqtt_user, mqtt_password);
    mqttClient.setId("ESP32_RA_Transmitter");
    mqttClient.onMessage(onMqttMessage);

    mqttReconnect();

    Serial.println("ESP32 RA Transmitter is ready.");
}

// --- Main Loop ---
void loop() {
    // 1. Maintain MQTT connection
    if (!mqttClient.connected()) {
        mqttReconnect();
    }
    mqttClient.poll();

    // 2. Check for incoming LoRa messages (ACKs, Status Updates)
    int packetSize = LoRa.parsePacket();
    if (packetSize) {
        String receivedMessage = "";
        while (LoRa.available()) {
            receivedMessage += (char)LoRa.read();
        }
        receivedMessage.trim();
        if (receivedMessage.length() > 0) {
            handleReceivedLoRaMessage(receivedMessage);
        }
    }

    // 3. Check resend queue periodically (non-blocking timer)
    unsigned long currentTime = millis();
    if (currentTime - lastTimerCheck >= timerInterval) {
        lastTimerCheck = currentTime;
        sendCounter++;
        if (sendCounter >= effectiveIntervalSeconds) {
            checkResendQueue();
            sendCounter = 0;
        }
    }

    // 4. Handle Serial Monitor Input (for debugging/testing)
    if (Serial.available() > 0) {
        String command = Serial.readStringUntil('\n');
        command.trim();
        Serial.print("Received from Serial Monitor: "); Serial.println(command);

        JsonDocument testDoc;
        DeserializationError testErr = deserializeJson(testDoc, command);

        if (!testErr) {
            if (testDoc.is<JsonObject>() && !testDoc["targetDevice"].isNull()) {
                String message_id = String(random(10000, 99999));
                JsonDocument loraDoc;
                String targetDeviceString = testDoc["targetDevice"].as<String>();
                loraDoc["message_id"] = message_id;
                loraDoc["targetDevice"] = targetDeviceString;
                JsonObjectConst obj = testDoc.as<JsonObjectConst>();
                for (JsonObjectConst::iterator it = obj.begin(); it != obj.end(); ++it) {
                    if (String(it->key().c_str()) != "targetDevice") {
                        loraDoc[it->key()] = it->value();
                    }
                }
                String jsonMessage;
                serializeJson(loraDoc, jsonMessage);
                sendMessage(jsonMessage, message_id, targetDeviceString);
            } else {
                if (!testDoc.is<JsonObject>()) {
                    Serial.println("Error: Serial Monitor input is not a JSON object.");
                } else {
                    Serial.println("Error: JSON from Serial Monitor must contain 'targetDevice'.");
                }
            }
        } else {
            Serial.println("Invalid JSON format received via Serial Monitor.");
            Serial.print("Parsing error: "); Serial.println(testErr.c_str());
        }
    }

    delay(5);
}

// --- WiFi and MQTT Functions ---
void connectToWiFi() {
    Serial.print("Initializing WiFi...");
    WiFi.mode(WIFI_STA);
    WiFi.disconnect(true);
    delay(100);
    Serial.print("Connecting to WiFi...");
    WiFi.begin(ssid, password);
    int retryCount = 0;
    while (WiFi.status() != WL_CONNECTED) {
        delay(1000);
        Serial.print(".");
        retryCount++;
        if (retryCount > 30) {
            Serial.println("\nFailed to connect to WiFi. Restarting...");
            ESP.restart();
        }
    }
    Serial.println("\nConnected to WiFi!");
    Serial.print("IP Address: "); Serial.println(WiFi.localIP());
}

void mqttReconnect() {
    while (!mqttClient.connected()) {
        Serial.print("Attempting MQTT connection...");
        if (mqttClient.connect(mqtt_server, 1883)) {
            Serial.println("connected");
            mqttClient.subscribe(mqtt_topic_subscribe);
            Serial.print("Subscribed to: "); Serial.println(mqtt_topic_subscribe);
        } else {
            Serial.print("failed, rc=");
            Serial.print(mqttClient.connectError());
            Serial.println(" trying again in 5 seconds");
            delay(5000);
        }
    }
}

// MQTT message handler
void onMqttMessage(int messageSize) {
    String topic = mqttClient.messageTopic();
    String messageReceived;
    while (mqttClient.available()) {
        messageReceived += (char)mqttClient.read();
    }
    Serial.println("--- MQTT Callback ---");
    Serial.println("Topic: " + topic);
    Serial.println("Raw Payload (" + String(messageSize) + " bytes): " + messageReceived);

    if (topic == mqtt_topic_subscribe) {
        JsonDocument doc;
        Serial.println("Attempting to parse MQTT JSON...");
        DeserializationError mqttError = deserializeJson(doc, messageReceived);

        if (mqttError == DeserializationError::NoMemory) {
            Serial.println("!!! Failed to allocate memory for MQTT JSON parsing.");
            return;
        } else if (mqttError) {
            Serial.print("!!! Failed to parse JSON from MQTT: ");
            Serial.println(mqttError.c_str());
            Serial.println("Payload was: " + messageReceived);
            return;
        }
        Serial.println("MQTT JSON Parsed Successfully.");

        if (doc["targetDevice"].isNull()) {
            Serial.println("!!! MQTT message missing 'targetDevice'. Cannot route over LoRa. Ignoring message.");
            return;
        }
        String targetDevice = doc["targetDevice"].as<String>();
        Serial.println("Target Device for LoRa: " + targetDevice);

        String message_id = String(random(10000, 99999));
        doc["message_id"] = message_id;
        Serial.println("Added message_id for LoRa transmission: " + message_id);

        String jsonMessageToSendViaLoRa;
        Serial.println("Serializing document for LoRa transmission...");
        size_t jsonSize = serializeJson(doc, jsonMessageToSendViaLoRa);
        Serial.println("Serialized LoRa JSON (" + String(jsonSize) + " bytes): " + jsonMessageToSendViaLoRa);

        sendMessage(jsonMessageToSendViaLoRa, message_id, targetDevice);
    } else {
        Serial.println("Message received on incorrect topic (not the configured control topic). Ignoring.");
    }
    Serial.println("--- End MQTT Callback ---");
}

// --- LoRa Communication Functions ---
void sendMessage(String jsonMessageToSendViaLoRa, String message_id, String targetDevice) {
    JsonDocument doc;
    DeserializationError error = deserializeJson(doc, jsonMessageToSendViaLoRa);
    if (error) {
        Serial.print("sendMessage: Invalid JSON format, cannot send via LoRa: ");
        Serial.println(error.c_str());
        return;
    }
    if (doc["message_id"].isNull() || doc["targetDevice"].isNull()) {
        Serial.println("sendMessage: JSON missing essential 'message_id' or 'targetDevice'. Cannot send.");
        return;
    }

    auto it = std::find_if(resendQueue.begin(), resendQueue.end(),
        [targetDevice](const Message& msg_in_queue) {
            return msg_in_queue.targetDevice == targetDevice;
        });

    if (it != resendQueue.end()) {
        Serial.println("Found existing message for target '" + targetDevice + "' (ID: " + it->message_id + ") in queue. Removing it before sending new message.");
        resendQueue.erase(it);
    }

    Serial.print("Sending LoRa packet to '" + targetDevice + "' (ID: " + message_id + ", Queue Size: " + String(resendQueue.size()) + "): ");
    Serial.println(jsonMessageToSendViaLoRa);
    LoRa.beginPacket();
    LoRa.print(jsonMessageToSendViaLoRa);
    int success = LoRa.endPacket();
    if(success){
        Serial.println("LoRa packet sent initially.");
    } else {
        Serial.println("LoRa packet initial send FAILED.");
    }

    Message newMsg = {jsonMessageToSendViaLoRa, message_id, targetDevice, millis(), 0};
    resendQueue.push_back(newMsg);
    Serial.println("Added message ID " + message_id + " for target '" + targetDevice + "' to queue (" + String(resendQueue.size()) + " items).");
}

void checkResendQueue() {
    unsigned long currentMillis = millis();
    unsigned long resendIntervalMillis = effectiveIntervalSeconds * 1000UL;

    if (resendQueue.empty()) {
        return;
    }

    for (auto it = resendQueue.begin(); it != resendQueue.end(); ) {
        if (it->retryCount >= resend_times) {
            Serial.println("Message ID " + it->message_id + " for target '" + it->targetDevice + "' failed after max retries (" + String(resend_times) + "). Removing from queue.");
            it = resendQueue.erase(it);
            continue;
        }

        if (currentMillis - it->lastSentTime >= resendIntervalMillis) {
            Serial.print("Resending message ID " + it->message_id + " for target '" + it->targetDevice + "' (Attempt " + String(it->retryCount + 1) + "/" + String(resend_times) + "): " + it->message);
            LoRa.beginPacket();
            LoRa.print(it->message);
            int success = LoRa.endPacket();
            if(success){
                Serial.println(" ...Resent Successfully.");
            } else {
                Serial.println(" ...Resend FAILED.");
            }
            it->lastSentTime = currentMillis;
            it->retryCount++;
            ++it;
        } else {
            ++it;
        }
    }
}

// Handles messages received VIA LORA from the receiver devices (ACKs, Status Updates, Discovery Config requests)
// Handles messages received VIA LORA from the receiver devices (ACKs, Status Updates, Discovery Config requests)
void handleReceivedLoRaMessage(String receivedMessage) {
    Serial.print(">>> Transmitter Received LoRa: ");
    Serial.println(receivedMessage);

    // Use StaticJsonDocument for parsing incoming messages. 512 bytes should be enough for ACK, Status, and compact Config.
    StaticJsonDocument<512> doc;
    DeserializationError error = deserializeJson(doc, receivedMessage); // Parse the received JSON
    if (error) {
        Serial.print(">>> Transmitter: Failed to parse received LoRa JSON: ");
        Serial.println(error.c_str());
        return; // Exit if JSON is invalid
    }

    // --- Handle ACK from Receiver ---
    // An ACK should contain "message":"acknowledged" and the original "message_id" it is acknowledging.
    // It should also contain the "device" field so we know which device sent the ACK.
    if (!doc["message"].isNull() && doc["message"] == "acknowledged" && !doc["message_id"].isNull() && !doc["device"].isNull()) {
        String ackMsgId = doc["message_id"].as<String>();
        String acknowledgingDevice = doc["device"].as<String>();
        Serial.printf(">>> Transmitter: Received ACK for message_id [%s] from device [%s].\n", ackMsgId.c_str(), acknowledgingDevice.c_str());

        // Find the matching message in the resendQueue based on the received message_id.
        // We need to find the message *we sent* that corresponds to this ACK.
        // Capture ackMsgId by value in the lambda for safe use.
        auto it = std::find_if(resendQueue.begin(), resendQueue.end(),
                               [ackMsgId](const Message& msg_in_queue) {
            // Compare the message_id from the received ACK with the message_id in the queue.
            return msg_in_queue.message_id == ackMsgId;
        });
        if (it != resendQueue.end()) {
            // Found the message that was acknowledged, remove it from the queue.
            Serial.println(">>> Transmitter: Found matching message (Target: " + it->targetDevice + ") in queue. Erasing.");
            resendQueue.erase(it); // Remove the acknowledged message
        } else {
            // This could happen if the ACK arrived after the message was already removed
            // (e.g., manually cleared, max retries reached, or a duplicate ACK).
            Serial.println(">>> Transmitter: ACK received, but no matching message found in queue for message_id " + ackMsgId + " (might have timed out or already ACKed).");
        }
        return; // Processed an ACK, exit the function
    }

    // --- Handle compact Home Assistant discovery config message from Receiver ---
    // The receiver sends this to request the gateway to set up HA discovery for one of its sensors/binary_sensors.
    // It contains "b", "s", and "n". It also includes "message_id".
    // It *does not* contain a top-level "device" key like status updates do (this is how we differentiate it).
    // ADDED: Check for 'n' (device name) to further confirm it's a config request.
    if (doc.containsKey("b") && doc.containsKey("s") && doc.containsKey("n") && doc.containsKey("message_id") && !doc.containsKey("device")) {
        Serial.println(">>> Transmitter: Received LoRa Home Assistant discovery config request (compact format).");

        // Extract necessary information from the compact message
        bool isBinary = doc["b"];                      // 1 for binary sensor, 0 for regular sensor
        String sensorType = doc["s"].as<String>();     // e.g., "motion", "illuminance", "voltage", "lightStatus"
        String deviceName = doc["n"].as<String>();     // e.g., "porch_light" (unique device identifier) from the receiver's perspective
        String receivedMsgId = doc["message_id"].as<String>(); // Message ID from the received config request

        // Validate essential fields are not empty
        if (deviceName.isEmpty() || sensorType.isEmpty() || receivedMsgId.isEmpty()) {
            Serial.println(">>> Transmitter: Received invalid discovery config (empty device/sensor/message_id/name). Ignoring.");
            // Optionally send a NACK here if the protocol supported it.
            return; // Ignore invalid config messages
        }

        // --- Construct components for the full Home Assistant discovery payload ---
        // unique_id: This is how HA tracks the entity long-term. Should be descriptive and unique across all devices.
        String unique_id = deviceName + "_" + sensorType;

        // object_id: This is the part of the topic. Making it shorter/different from unique_id seems to fix mutation.
        // ADDED: Logic to create a specific object_id based on sensor type for the topic.
        String object_id_suffix = sensorType.substring(0, min((unsigned int)sensorType.length(), 4u)); // Default suffix (first 4 chars)
        if (sensorType == "lightStatus") {
             object_id_suffix = "lstat"; // Specific short suffix for lightStatus binary sensor
        } else if (sensorType == "motion") {
             object_id_suffix = "motion"; // Specific suffix for motion binary sensor
        } else if (sensorType == "illuminance") {
             object_id_suffix = "ill"; // Specific suffix for illuminance sensor
        } else if (sensorType == "voltage") {
             object_id_suffix = "volt"; // Specific suffix for voltage sensor
        }
        // Add else if for other sensor types as needed based on your receiver config

        String object_id_topic = deviceName + "_" + object_id_suffix; // Use device name + suffix for topic


        String device_class = sensorType;                     // Default device class (can be refined below)
        String state_topic = mqtt_topic_publish;              // All sensor states will be published to this common topic
        String value_template = "{{ value_json." + sensorType + " }}"; // Default template to extract the specific sensor value
        String payload_on = "";                               // Default for binary sensors
        String payload_off = "";                              // Default for binary sensors
        String unit = "";                                     // Default unit of measurement
        String state_class = "";                              // Default state_class (optional, good for numerical sensors)

        // Refine device_class, unit, value_template, state_class based on specific sensor types
        if (isBinary) { // This block handles {"b": 1, ...}
            // Specific binary sensor types based on 's' (sensorType)
            if (sensorType == "motion") {
                device_class = "motion"; // Home Assistant standard device class
                payload_on = "1";        // Value in the state JSON indicating ON (from the receiver)
                payload_off = "0";       // Value in the state JSON indicating OFF (from the receiver)
                 value_template = "{{ value_json.motion }}"; // Ensure template matches expected key in state payload
            } else if (sensorType == "lightStatus") { // <--- Handles the new lightStatus binary sensor config
                device_class = "light"; // Use standard 'light' device class for binary light state
                payload_on = "on";        // The string value in your state payload for ON ("on")
                payload_off = "off";       // The string value in your state payload for OFF ("off")
                value_template = "{{ value_json.lightStatus }}"; // Template matches key in state payload
                unit = ""; // No unit for binary sensor
                state_class = ""; // No state class for binary sensor
            }
             // ADD other binary sensor types as needed
        } else { // This block handles {"b": 0, ...} (Regular sensors)
            // Specific regular sensor types based on 's' (sensorType)
            if (sensorType == "illuminance") {
                device_class = "illuminance"; // Home Assistant standard
                unit = "lx";                 // Unit of measurement
                state_class = "measurement";  // For continuous numerical data
                 value_template = "{{ value_json.lux }}"; // Ensure template matches expected key in state payload
            } else if (sensorType == "voltage") {
                device_class = "voltage";     // Home Assistant standard
                unit = "V";                   // Unit of measurement
                state_class = "measurement";  // For continuous numerical data
                 value_template = "{{ value_json.voltage }}"; // Ensure template matches expected key in state payload
            } else if (sensorType == "lightMode") { // Keep lightMode here (it's not binary)
                device_class = "enum";
                unit = "";
                state_class = "";
                value_template = "{{ value_json.lightMode }}"; // Template matches the key in state payload
            }
            // ADD other regular sensor types as needed
        }

        // --- Build the FULL Home Assistant discovery payload ---
        // Use a StaticJsonDocument with a size large enough for the full HA payload, including the device block.
        // 768 bytes is an increased size, which should be sufficient for most standard sensor/binary_sensor payloads with a device block.
        StaticJsonDocument<768> ha_payload;

        // Add entity-specific configuration
        ha_payload["unique_id"] = unique_id;             // Unique ID for this specific entity (sensor or binary_sensor)
        ha_payload["device_class"] = device_class;       // Home Assistant device class
        ha_payload["name"] = deviceName + " " + sensorType; // Friendly name for the entity (e.g., "porch_light motion")
        ha_payload["state_topic"] = state_topic;         // Topic where this entity's state updates will be published
        ha_payload["value_template"] = value_template;   // How to extract the state value from the state_topic JSON

        // Add optional fields if they are applicable/set
        if (!unit.isEmpty()) {
            ha_payload["unit_of_measurement"] = unit;
        }
        if (!state_class.isEmpty()) {
             ha_payload["state_class"] = state_class; // Add state_class for numerical sensors
        }
        if (isBinary) { // Only add payload_on/off for binary sensors
            ha_payload["payload_on"] = payload_on;
            ha_payload["payload_off"] = payload_off;
        }
        // Add other common options like "icon", "entity_category" if desired

        // Add the device block - this groups entities belonging to the same physical device in Home Assistant
        // This block should be identical for all entities belonging to the same physical device (porch_light)
        JsonObject dev = ha_payload.createNestedObject("device");
        dev["identifiers"][0] = deviceName; // Unique identifier for the physical device (e.g., "porch_light")
        dev["name"] = deviceName;         // Name of the physical device in Home Assistant
        dev["model"] = "LoRa Sensor";     // Device model (customize as appropriate for your devices)
        dev["manufacturer"] = "Custom";   // Manufacturer (customize as appropriate)
        // You can add more device info here if available/desired (e.g., sw_version, hw_version, via_device - via_device could link it to the gateway)

        // --- Serialize the full payload to a character buffer ---
        // The buffer size must be large enough to hold the serialized JSON string.
        // Using 800 bytes as a safer buffer size corresponding to the 768 byte doc.
        char payloadBuffer[800];
        size_t jsonSize = serializeJson(ha_payload, payloadBuffer);

        if (jsonSize == 0 || jsonSize >= sizeof(payloadBuffer)) {
            Serial.println(">>> Transmitter: Failed to serialize HA discovery payload to buffer (size 0 or buffer too small). Cannot publish config.");
            // If serialization fails, we cannot publish or send ACK based on publish success.
            return; // Exit handler
        }
        // ensure null termination (serializeJson should do this, but defensive programming doesn't hurt)
        payloadBuffer[jsonSize] = '\0';

        // --- Determine the correct discovery topic for this entity using the modified object_id ---
        String config_topic = "homeassistant/";
        config_topic += isBinary ? "binary_sensor/" : "sensor/"; // Use binary_sensor or sensor topic prefix
        config_topic += object_id_topic;                         // **Use the modified object_id here**
        config_topic += "/config";                               // The standard HA discovery config suffix

        // --- Publish the serialized payload using the robust publish function ---
        // It is important that the discovery config is published reliably AND RETAINED.
        Serial.printf(
            ">>> Transmitter: Attempting robust MQTT publish for HA config to topic [%s] (object_id: %s, unique_id: %s, %u bytes)\n",
            config_topic.c_str(),
            object_id_topic.c_str(), // Print object_id for debugging
            unique_id.c_str(),       // Print unique_id for debugging
            jsonSize
        );

        // payloadBuffer is a char array, cast it to uint8_t* for publishRobustly
        // Publish with retain: true so Home Assistant finds it on restart
        bool pubSuccess = publishRobustly(
            config_topic.c_str(),       // The calculated config topic
            ha_payload,                 // **Pass the JsonDocument directly**
            true,                       // RETAIN this message on the broker
            3                           // max_attempts
        );

        // --- Send ACK back over LoRa to the receiver ONLY IF the MQTT publish of the config was successful ---
        // The receiver requested this config and is waiting for an ACK for this specific message ID.
        // This ACK tells the receiver the gateway successfully handled the config request and published to MQTT.
        if (pubSuccess) {
            StaticJsonDocument<128> ackDoc; // Small doc for the ACK message
            ackDoc["message_id"] = receivedMsgId; // Use the message_id from the original received config message
            ackDoc["message"] = "acknowledged";   // Indicate successful processing by the gateway
            // Optionally add a gateway identifier to the ACK
            // ackDoc["gateway"] = "ESP32_RA_Transmitter"; // Let the receiver know who ACKed

            String ackMessage;
            serializeJson(ackDoc, ackMessage); // Serialize the ACK message

            Serial.printf(
                ">>> Transmitter: MQTT Publish successful for config %s. Sending LoRa ACK for message_id %s: %s\n",
                unique_id.c_str(), // Which config was successful
                receivedMsgId.c_str(),
                ackMessage.c_str()
            );
            LoRa.beginPacket(); // Start LoRa packet for ACK
            LoRa.print(ackMessage); // Write ACK JSON string
            LoRa.endPacket(); // Finish and send packet (blocking)
        } else {
            // If robust publish failed after all attempts, we do NOT send an ACK.
            // The receiver's waitForAck for this config message will time out, and it should retry sending the config message later.
            Serial.printf(
                ">>> Transmitter: MQTT Publish FAILED for config %s after attempts for message_id %s. NOT sending ACK. Receiver should retry.\n",
                unique_id.c_str(),
                receivedMsgId.c_str()
            );
            // Consider adding logging or notification here for persistent failure to publish config
        }

        return; // Processed a discovery config message, exit the function
    }

    // --- Handle General Status Update from Receiver ---
    // If the message was not an ACK and not a discovery config, check if it's a general status update.
    // A status update from the receiver should contain at least a "device" identifier
    // and the actual sensor readings or states (e.g., "motion", "illuminance", "voltage", "lightStatus").
    // ADDED: Ensure it does NOT contain 'b', 's', 'n', 'message_id' to avoid confusion with compact config.
    if (!doc["device"].isNull() && !doc.containsKey("b") && !doc.containsKey("s") && !doc.containsKey("n")) {
        // Assuming the entire received message is the status payload Home Assistant expects on the state topic.
        // The receiver's 'sendStatusUpdate' function should format this correctly.
        // We will publish the *entire received JSON string* directly to the state topic.

        String statusPayload = receivedMessage; // Use the original received string
        String reportingDevice = doc["device"].as<String>();

        Serial.print(
            ">>> Transmitter: Received LoRa Status Update from device [" + reportingDevice + "]. Publishing entire received JSON to MQTT topic ["
        );
        Serial.print(mqtt_topic_publish);
        Serial.print("]: ");
        // Republish the entire received JSON string directly to the state topic
        Serial.println(statusPayload);

        // Publish the status message. Use QoS 0 and RETAIN false for state updates.
        // Use your updated publishRobustly for consistency, though QoS 0/not retained is less critical.
bool pubSuccess = publishRobustly(
            mqtt_topic_publish,         // The state topic
            doc,                        // **Pass the parsed JsonDocument**
            false,                      // DO NOT RETAIN state updates
            1                           // max_attempts
        );

        if(!pubSuccess){
             Serial.println("!!! Transmitter: Failed to publish status update to MQTT.");
        }
        return; // Processed a status update, exit the function
    }

    // --- Handle settings update from Receiver (THIS SHOULD NOT HAPPEN IN THIS DESIGN) ---
    // In this architecture, settings are PUSHED from Home Assistant via the MQTT callback,
    // not PULLED by the receiver requesting them from the gateway via LoRa.
    // The gateway should NOT contain logic to generate settings in response to a "request_settings"
    // action from the receiver. That logic belongs in Home Assistant automations or Node-RED flows.
    // If your receiver sends {"action":"request_settings", "device":"..."} and your previous gateway
    // code had logic here to respond with settings, that logic should be removed or commented out
    // to enforce the intended architecture.
    /*
    // Example of logic to REMOVE if present in your code:
    if (!doc["action"].isNull() && doc["action"] == "request_settings" && !doc["device"].isNull()) {
        Serial.println(">>> Transmitter: Received request_settings from device. IGNORING as settings are pushed from MQTT callback.");
        // DO NOT generate and send settings here. The process is: HA (MQTT) -> Gateway (callback) -> LoRa (sendMessage).
        return; // Ignore this message type based on the intended design
    }
    */

    // --- If message was none of the above recognized types ---
    Serial.println(
        ">>> Transmitter: Received unhandled LoRa message format."
    );
    Serial.println("Raw message: " + receivedMessage); // Print unhandled message
    // If you receive this frequently, you might need to analyze the receivedMessage
    // content to understand what the receiver is sending that isn't being matched.
    // Optionally send an error ACK back to the receiver if your protocol supports it.
}



// --- Robust MQTT Publish Function ---
// Modified to accept JsonDocument directly, a 'retain' flag, and max_attempts
bool publishRobustly(const char* topic, const JsonDocument& doc, bool retain, int max_attempts) {
    int attempts = 0;
    bool published = false;

    // --- CALCULATE JSON SIZE ONCE, AT THE BEGINNING ---
    size_t jsonSize = measureJson(doc);

    // --- Debug print, now jsonSize is available ---
    Serial.printf("publishRobustly: Attempting to publish to %s (retain=%d, %u bytes)...\n", topic, retain, jsonSize);

    // --- Debug print of payload content (using the 'doc' JsonDocument) ---
    Serial.println("--- Payload JSON Content (before beginMessage) ---");
    serializeJson(doc, Serial); // Print to Serial
    Serial.println(); // Add a newline after the JSON output
    Serial.println("------------------------------------------");
    // -------------------------------------------------


    while (attempts < max_attempts && !published) {
        if (mqttClient.connected()) {
            Serial.printf("publishRobustly: Attempt %d/%d. MQTT IS connected. Sending...\n", attempts + 1, max_attempts);

            // --- NEW PUBLISHING CODE (using jsonSize from outside the loop) ---
            // Begin message with the exact size announced upfront
            // (Assumes QoS 0 as per your typical use)
            mqttClient.beginMessage(topic, jsonSize, retain);

            // Stream the serialized JSON directly to the MQTT client
            size_t bytesWritten = serializeJson(doc, mqttClient);

            // End the message
            mqttClient.endMessage();

            // Check if all bytes were written (should match the measured size)
            if (bytesWritten == jsonSize) {
                published = true;
                Serial.printf("publishRobustly: serializeJson and endMessage() success. %u of %u bytes written.\n", bytesWritten, jsonSize);
            } else {
                // This indicates a failure during serialization or streaming
                Serial.printf("publishRobustly: WARNING: serializeJson only wrote %u of %u bytes. Publish failed.\n", bytesWritten, jsonSize);
                published = false; // Mark as failed
            }
            // --- END OF NEW PUBLISHING CODE ---

        } else {
            Serial.println("publishRobustly: MQTT NOT connected. Trying to reconnect...");
            mqttReconnect(); // Use your existing reconnect logic
            delay(500); // Small delay after reconnect attempt
            // Note: attempt counter increment is handled by the loop condition implicitly
        }
        // If not published in this attempt (either disconnected or write failed)
        if (!published) {
             if (!mqttClient.connected()){
                  // If connect failed, the delay inside mqttReconnect already happened.
                  // The attempt counter was effectively incremented by looping.
                   Serial.println("publishRobustly: Publish attempt failed (MQTT not connected). Will retry in next loop iteration.");
             } else {
                 // If connected but write failed
                 Serial.println("publishRobustly: Publish attempt failed (Write/Stream error). Waiting before next retry...");
                 delay(1000); // Wait a bit longer between failed attempts if connected but write failed
                 attempts++; // Manually increment attempt count if connected but write failed
             }
        }
    }
    Serial.printf("publishRobustly: Finished. Final published status = %d\n", published ? 1 : 0);
    return published;
}