#include "mqtt.hpp"
#include "config.hpp"
#include <PubSubClient.h>
#include "logger.h"
#include <ArduinoJson.h>

extern logging::Logger logger;
extern Config config;

Mqtt::Mqtt(WiFiClient& wifiClient) : mqttClient(wifiClient) {
    deviceId = "water_temp_receiver";
    baseTopic = "homeassistant/sensor/water_temp";
}

void Mqtt::mqttSetup()
{
    String message = "MqttServer: " + config.mqttConfig.server + ", MqttPort: " + config.mqttConfig.port;
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", message.c_str());
    if (config.mqttConfig.server.length() > 0)
    {
        String message = "MqttServer: " + config.mqttConfig.server + ", MqttPort: " + config.mqttConfig.port;
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", message.c_str());
        IPAddress mqttServerIP;
        if(mqttServerIP.fromString(config.mqttConfig.server)){
            
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", mqttServerIP.toString().c_str());
            mqttClient.setServer(mqttServerIP, config.mqttConfig.port);
        }else{
            logger.log(logging::LoggerLevel::LOGGER_LEVEL_ERROR, "MQTT", "MQTT Invalid IP address format %s", config.mqttConfig.server);
        }
        mqttClient.setCallback(callback);
    }
    else
    {
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_ERROR, "MQTT", "MQTT config is not set");
    }
}

void Mqtt::mqttRun()
{
    delay(100);
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "MQTT RUN");
    if (WiFi.status() == WL_CONNECTED && !mqttClient.loop()) {
        mqttReconnect();
    } 
}

void Mqtt::mqttReconnect()
{
    int retryCount = 0;
    
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Attempting MQTT connection...");
    // Loop until we're reconnected
    
    while (!mqttClient.connected() && retryCount < 10)
    {
        // Attempt to connect
        String mqttId = config.mqttConfig.id;
        if (mqttId.isEmpty())
        {
            mqttId = "ESP32-";
            mqttId += String(random(0xffff), HEX);
            
            logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Set ID: %s", mqttId );
            config.setMqttId(mqttId);
            config.writeData();
        }

        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT CONFIG", "Server: %s", config.mqttConfig.server.c_str());
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT CONFIG", "Port: %d", config.mqttConfig.port);
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT CONFIG", "Username: %s", config.mqttConfig.username.c_str());
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT CONFIG", "ID: %s", config.mqttConfig.id.c_str());
  
        if (mqttClient.connect(mqttId.c_str(), config.mqttConfig.username.c_str(), config.mqttConfig.password.c_str()))
        {
            logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "connected");
            
            // Setup Home Assistant Discovery on successful connection
            setupHomeAssistantDiscovery();
            
            String topic = config.mqttConfig.topic;
            if (topic.isEmpty())
            {
                topic = config.mqttConfig.id;
            }
            mqttClient.subscribe(topic.c_str());
            String message = "Subscribe to: " + topic;
            logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", message.c_str());
        }
        else
        {
            logger.log(logging::LoggerLevel::LOGGER_LEVEL_WARN, "MQTT", "Not connected to mqtt, try again in 5 s.");
            // Wait 5 seconds before retrying
            delay(5000);
        }
    }
    // Check if we're connected after the loop
    if (!mqttClient.connected()) {
        // Connection failed after maximum retries
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_ERROR, "MQTT", "Failed to connect to MQTT broker after 10 retries");
    }
}

void Mqtt::setupHomeAssistantDiscovery()
{
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Setting up Home Assistant MQTT Discovery");
    
    // Temperature sensor discovery
    {
        JsonDocument doc;
        doc["name"] = "Water Temperature";
        doc["unique_id"] = "water_temp_temperature";
        doc["unit_of_measurement"] = "°C";
        doc["device_class"] = "temperature";
        doc["state_topic"] = baseTopic + "/temperature";
        doc["availability_topic"] = baseTopic + "/availability";
        
        JsonObject device = doc["device"].to<JsonObject>();
        device["identifiers"][0] = deviceId;
        device["name"] = "Water Temperature Sensor";
        device["model"] = "Water Temp v1";
        device["manufacturer"] = "DIY";
        
        String payload;
        serializeJson(doc, payload);
        String topic = baseTopic + "/temperature/config";
        mqttPublish(topic.c_str(), payload.c_str());
    }
    
    // RSSI sensor discovery
    {
        JsonDocument doc;
        doc["name"] = "LoRa RSSI";
        doc["unique_id"] = "water_temp_rssi";
        doc["unit_of_measurement"] = "dBm";
        doc["device_class"] = "signal_strength";
        doc["state_topic"] = baseTopic + "/rssi";
        doc["availability_topic"] = baseTopic + "/availability";
        
        JsonObject device = doc["device"].to<JsonObject>();
        device["identifiers"][0] = deviceId;
        
        String payload;
        serializeJson(doc, payload);
        String topic = baseTopic + "/rssi/config";
        mqttPublish(topic.c_str(), payload.c_str());
    }
    
    // Battery level discovery
    {
        JsonDocument doc;
        doc["name"] = "Sender Battery";
        doc["unique_id"] = "water_temp_battery";
        doc["unit_of_measurement"] = "%";
        doc["device_class"] = "battery";
        doc["state_topic"] = baseTopic + "/battery";
        doc["availability_topic"] = baseTopic + "/availability";
        
        JsonObject device = doc["device"].to<JsonObject>();
        device["identifiers"][0] = deviceId;
        
        String payload;
        serializeJson(doc, payload);
        String topic = baseTopic + "/battery/config";
        mqttPublish(topic.c_str(), payload.c_str());
    }
    
    // Set availability to online
    String availTopic = baseTopic + "/availability";
    mqttPublish(availTopic.c_str(), "online");
    
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Home Assistant Discovery setup complete");
}

void Mqtt::publishTemperature(float temp)
{
    String topic = baseTopic + "/temperature";
    String payload = String(temp, 2);
    mqttPublish(topic.c_str(), payload.c_str());
}

void Mqtt::publishRSSI(int rssi)
{
    String topic = baseTopic + "/rssi";
    String payload = String(rssi);
    mqttPublish(topic.c_str(), payload.c_str());
}

void Mqtt::publishSensorData(float temp, int rssi, int battery)
{
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Publishing sensor data - Temp: %.2f, RSSI: %d, Battery: %d", temp, rssi, battery);
    
    publishTemperature(temp);
    publishRSSI(rssi);
    
    String topic = baseTopic + "/battery";
    String payload = String(battery);
    mqttPublish(topic.c_str(), payload.c_str());
    
    // Also publish combined JSON
    {
        JsonDocument doc;
        doc["temperature"] = serialized(String(temp, 2));
        doc["rssi"] = rssi;
        doc["battery"] = battery;
        doc["timestamp"] = millis();
        
        String jsonPayload;
        serializeJson(doc, jsonPayload);
        String jsonTopic = baseTopic + "/state";
        mqttPublish(jsonTopic.c_str(), jsonPayload.c_str());
    }
}

void Mqtt::callback(char *topic, uint8_t *payload, unsigned int length)
// Not implemented yet. Just for debug
{
    // Create spot in memory for message
    char message[length + 1];
    // Write payload to String
    strncpy(message, (char *)payload, length);
    // Nullify last character to eliminate garbage at end
    message[length] = '\0';

    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Received: %s", message);
    logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "From: %s", topic);
}

void Mqtt::mqttPublish(const char* topic, const char* payload){
     if (mqttClient.connected()) {
        mqttClient.publish(topic, payload);
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_INFO, "MQTT", "Published message to %s: %s", topic, payload);
    } else {
        logger.log(logging::LoggerLevel::LOGGER_LEVEL_ERROR, "MQTT", "Failed to publish message. MQTT client not connected.");
    }
}
