#include <ArduinoJson.h>
#include "light.h"
#include "mqtt.h"

constexpr uint16_t configMsgSize = 1024;
constexpr uint16_t statusMsgSize = 128;
constexpr uint8_t minPwmValue = 0;
constexpr double gammaCorrection = 2.4;
constexpr uint32_t maxPwmUi = 255;

const struct {
    String available = "online";
    String notAvailable = "offline";
} Availability;


Light::Light(Pin* pinR, Pin* pinG, Pin* pinB, Mqtt* mqttClient, std::string uniqueId)
    : pinR(pinR), pinG(pinG), pinB(pinB), pinCW(nullptr), pinWW(nullptr), mqttClient(mqttClient) {
    lightInfo.uniqueId = uniqueId;
    lightType = LightType::rgb;
    uint8_t bits = pinR->getLedResolutionBits();
    maxPwm = (bits >= 1 && bits <= 31) ? ((1u << bits) - 1u) : 255u;
    publishInitialState();
    subscribeToMqttTopics();
}

Light::Light(Pin* pinR, Pin* pinG, Pin* pinB, Pin* pinCW, Pin* pinWW, Mqtt* mqttClient, std::string uniqueId)
    : pinR(pinR), pinG(pinG), pinB(pinB), pinCW(pinCW), pinWW(pinWW), mqttClient(mqttClient) {
    lightInfo.uniqueId = uniqueId;
    lightType = LightType::rgbww;
    uint8_t bits = pinR->getLedResolutionBits();
    maxPwm = (bits >= 1 && bits <= 31) ? ((1u << bits) - 1u) : 255u;
    publishInitialState();
    subscribeToMqttTopics();
}

void Light::publishInitialState() {
    // Publish the initial state of the light
    JsonDocument configInfo;
    JsonObject deviceInfo = configInfo["device"].to<JsonObject>();
    deviceInfo["name"] = this->deviceInfo.name;
    deviceInfo["model"] = this->deviceInfo.model;
    JsonArray identifiers = deviceInfo["identifiers"].to<JsonArray>();
    identifiers.add(this->deviceInfo.identifier + this->lightInfo.uniqueId);
    deviceInfo["sw_version"] = this->deviceInfo.swVersion;
    deviceInfo["manufacturer"] = this->deviceInfo.manufacturer;

    configInfo["unique_id"] = this->lightInfo.uniqueId;
    configInfo["name"] = this->lightInfo.name;
    configInfo["schema"] = "json";
    configInfo["json_attributes_topic"] = this->lightInfo.jsonAttributesTopic;
    configInfo["command_topic"] = this->lightInfo.commandTopic;
    configInfo["color_temp_kelvin"] = true;
    configInfo["max_kelvin"] = cwTempKelvin;
    configInfo["min_kelvin"] = wwTempKelvin;
    JsonArray availabilityInfo = configInfo["availability"].to<JsonArray>();
    JsonObject availabilityItem = availabilityInfo.add<JsonObject>();
    availabilityItem["topic"] = this->lightInfo.availabilityTopic;
    availabilityItem["value_template"] = this->lightInfo.availabilityTemplate;
    JsonArray supportedColorModes = configInfo["supported_color_modes"].to<JsonArray>();
    if (lightType == LightType::rgb) {
        supportedColorModes.add("rgb");
    } else if (lightType == LightType::rgbw) {
        supportedColorModes.add("rgbw");
    } else if (lightType == LightType::rgbww) {
        supportedColorModes.add("rgb");
        supportedColorModes.add("color_temp");
    } else if (lightType == LightType::colorTemperature) {
        supportedColorModes.add("color_temp");
    } else if (lightType == LightType::brightness) {
        supportedColorModes.add("brightness");
    } else {
        supportedColorModes.add("onoff");
    }
    configInfo["state_topic"] = this->lightInfo.stateTopic;

    std::string configJson;
    serializeJson(configInfo, configJson);
    mqttClient->publish(lightInfo.discoveryTopic, configJson);

    std::string stateJson;
    JsonDocument stateInfo;
    stateInfo["state"] = "OFF";
    stateInfo["brightness"] = maxPwmUi;
    brightness = maxPwmUi;
    JsonObject color = stateInfo["color"].to<JsonObject>();
    color["r"] = 0;
    r = 0;
    color["g"] = 0;
    g = 0;
    color["b"] = 0;
    b = 0;
    serializeJson(stateInfo, stateJson);
    std::string availabilityJson;
    JsonDocument availabilityInfoDoc;
    availabilityInfoDoc["availability"] = Availability.available;
    serializeJson(availabilityInfoDoc, availabilityJson);
    mqttClient->publish(lightInfo.stateTopic, stateJson);
    mqttClient->publish(lightInfo.availabilityTopic, availabilityJson);
}

void Light::operatePin() {
    auto clamp8 = [](int v)->uint8_t { return v < 0 ? 0 : (v > 255 ? 255 : v); };
    if (!isOn) {
        turnOff();
        return;
    }
    uint8_t r8 = clamp8(r);
    uint8_t g8 = clamp8(g);
    uint8_t b8 = clamp8(b);
    uint8_t cw8 = clamp8(cw);
    uint8_t ww8 = clamp8(ww);
    uint8_t br8 = clamp8(brightness);

    uint32_t rGamma = correctGamma(r8);
    uint32_t gGamma = correctGamma(g8);
    uint32_t bGamma = correctGamma(b8);
    uint32_t cwGamma = correctGamma(cw8);
    uint32_t wwGamma = correctGamma(ww8);
    uint32_t brGamma = correctGamma(br8);

    auto mixHw = [this](uint32_t cG, uint32_t bG) -> uint32_t {
        return static_cast<uint32_t>((static_cast<uint64_t>(cG) * bG + (maxPwm / 2)) / maxPwm);
    };

    uint32_t rSetpoint = mixHw(rGamma, brGamma);
    uint32_t gSetpoint = mixHw(gGamma, brGamma);
    uint32_t bSetpoint = mixHw(bGamma, brGamma);
    uint32_t cwSetpoint = 0;
    uint32_t wwSetpoint = 0;

    if (activeMode == ActiveMode::modeCct) {
        uint32_t peak = cwGamma > wwGamma ? cwGamma : wwGamma;
        if (peak > 0) {
            uint32_t cwGammaUp = static_cast<uint32_t>(
                (static_cast<uint64_t>(cwGamma) * maxPwm + (peak / 2)) / peak
            );
            uint32_t wwGammaUp = static_cast<uint32_t>(
                (static_cast<uint64_t>(wwGamma) * maxPwm + (peak / 2)) / peak
            );
            if (cwGammaUp > maxPwm) cwGammaUp = maxPwm;
            if (wwGammaUp > maxPwm) wwGammaUp = maxPwm;

            cwSetpoint = mixHw(cwGammaUp, brGamma);
            wwSetpoint = mixHw(wwGammaUp, brGamma);
        } else {
            cwSetpoint = wwSetpoint = 0;
        }
    } else {
        cwSetpoint = mixHw(cwGamma, brGamma);
        wwSetpoint = mixHw(wwGamma, brGamma);
    }

    if (pinR) pinR->setLedLevel(rSetpoint);
    if (pinG) pinG->setLedLevel(gSetpoint);
    if (pinB) pinB->setLedLevel(bSetpoint);
    if (pinCW) pinCW->setLedLevel(cwSetpoint);
    if (pinWW) pinWW->setLedLevel(wwSetpoint);
}

void Light::subscribeToMqttTopics() {
    mqttClient->subscribe(lightInfo.commandTopic, [this](uint8_t* payload, int length) {
        std::string command(reinterpret_cast<char*>(payload), length);
        handleCommand(command);
    });
}

void Light::handleCommand(const std::string& command) {
    Serial.println("Received command: " + String(command.c_str()));
    JsonDocument commandJson;
    deserializeJson(commandJson, command);
    if (commandJson.isNull()) {
        Serial.println("Invalid command JSON");
        return;
    }
    if (commandJson["state"].is<String>()) {
        std::string state = commandJson["state"].as<std::string>();
        if (state == "ON") {
            isOn = true;
        } else if (state == "OFF") {
            isOn = false;
        }
    }
    if (commandJson["brightness"].is<int>()) {
        brightness = commandJson["brightness"].as<int>();
    }
    if (commandJson["color"].is<JsonObject>()) {
        JsonObject color = commandJson["color"];
        r = color["r"] | maxPwmUi;
        g = color["g"] | maxPwmUi;
        b = color["b"] | maxPwmUi;
        cw = 0;
        ww = 0;
        activeMode = ActiveMode::modeRgb;
    }
    if (commandJson["color_temp"].is<int>()) {
        colorTemperature = commandJson["color_temp"].as<int>();
        applyKelvin(colorTemperature);
    }
    if (lightType == LightType::rgb || lightType == LightType::rgbw || lightType == LightType::rgbww)
    {
        operatePin();
        publishCurrentState();
    }
}

void Light::turnOff() {
    isOn = false;
    if (pinR != nullptr) pinR->setLedLevel(0);
    if (pinG != nullptr) pinG->setLedLevel(0);
    if (pinB != nullptr) pinB->setLedLevel(0);
    if (pinCW != nullptr) pinCW->setLedLevel(0);
    if (pinWW != nullptr) pinWW->setLedLevel(0);
}

void Light::publishCurrentState() {
    // Publish the current state of the light
    JsonDocument stateInfo;
    JsonDocument attributeInfo;
    stateInfo["state"] = isOn ? "ON" : "OFF";
    stateInfo["availability"] = Availability.available; // Current availability
    stateInfo["brightness"] = brightness;
    if (activeMode == ActiveMode::modeRgb) {
        JsonObject color = stateInfo["color"].to<JsonObject>();
        color["r"] = r;
        color["g"] = g;
        color["b"] = b;
        stateInfo["color_mode"] = "rgb";
    } else if (activeMode == ActiveMode::modeCct) {
        stateInfo["color_temp"] = colorTemperature;
        stateInfo["color_mode"] = "color_temp";
    }
    attributeInfo["pwmR"] = pinR->getLedLevel();
    attributeInfo["pwmG"] = pinG->getLedLevel();
    attributeInfo["pwmB"] = pinB->getLedLevel();
    if (pinCW != nullptr)
        attributeInfo["pwmCW"] = pinCW->getLedLevel();
    if (pinWW != nullptr)
        attributeInfo["pwmWW"] = pinWW->getLedLevel();

    std::string stateJson;
    serializeJson(stateInfo, stateJson);
    Serial.println("Publishing current state: " + String(stateJson.c_str()));
    mqttClient->publish(lightInfo.stateTopic, stateJson);
    std::string attributeJson;
    serializeJson(attributeInfo, attributeJson);
    Serial.println("Publishing current attributes: " + String(attributeJson.c_str()));
    mqttClient->publish(lightInfo.jsonAttributesTopic, attributeJson);
}

uint32_t Light::correctGamma(uint32_t originalPwm) {
    double normalized = static_cast<double>(originalPwm) / maxPwmUi;
    if (normalized <= 0.04045) {
        return static_cast<uint32_t>((normalized / 12.92) * maxPwm);
    } else {
        return static_cast<uint32_t>(pow((normalized + 0.055) / 1.055, gammaCorrection) * maxPwm);
    }
}

void Light::applyKelvin(uint32_t kelvin) {
    if (kelvin > cwTempKelvin) kelvin = cwTempKelvin;
    if (kelvin < wwTempKelvin) kelvin = wwTempKelvin;
    double tLin = static_cast<double>(kelvin - wwTempKelvin) / static_cast<double>(cwTempKelvin - wwTempKelvin);
    r = 0;
    g = 0;
    b = 0;
    cw = static_cast<uint8_t>(tLin * maxPwmUi);
    ww = static_cast<uint8_t>((1.0 - tLin) * maxPwmUi);
    activeMode = ActiveMode::modeCct;
}