#include "discovery.hpp"
#include <rapidjson/document.h>
#include <rapidjson/error/en.h>
#include <stinger/utils/format.hpp>
#include <stinger/utils/hash.hpp>
#include <iostream>
#include <algorithm>
#include <sstream>

{%for ns in config.cpp.namespace%}
namespace {{ns}} {
{%endfor-%}
namespace {{stinger.name|snake_case}} {

{%if stinger.properties|length > 0%}
bool InitialPropertyValues::isComplete() const
{
    {%-for prop_name, prop in stinger.properties.items() %}
    if (!{{prop_name | lowerCamelCase}}.has_value())
    {
        return false;
    }
    {%-endfor%}
    return true;
}
{%endif%}

bool InstanceInfo::isComplete() const
{
    if (!serviceId.has_value()) {
        return false;
    }
    {%for topic_param in config.topics.params %}
    if (!{{topic_param | lowerCamelCase}}.has_value()) {
        return false;
    }
    {%endfor%}
    {%- if stinger.properties|length > 0 %}
    if (!initial_property_values.isComplete()) {
        return false;
    }
    {%- endif %}
    return true;
}

void InstanceInfo::UpdateFromRapidJsonObject(const rapidjson::Value& jsonObj)
{
    if (jsonObj.HasMember("service_id") && jsonObj["service_id"].IsString())
    {
        serviceId = jsonObj["service_id"].GetString();
    }
    {%-for topic_param in config.topics.params %}
    if (jsonObj.HasMember("{{topic_param}}") && jsonObj["{{topic_param}}"].IsString())
    {
        {{topic_param|lowerCamelCase}} = jsonObj["{{topic_param}}"].GetString();
    }
    {%-endfor%}
};

void InstanceInfo::AddToRapidJsonObject(rapidjson::Value& parent, rapidjson::Document::AllocatorType& allocator) const
{
    parent.AddMember("interface_name", rapidjson::Value("{{stinger.name}}", allocator), allocator);
    parent.AddMember("interface_version", rapidjson::Value("{{stinger.version}}", allocator), allocator);
    if (serviceId.has_value())
    {
        parent.AddMember("service_id", rapidjson::Value(serviceId.value().c_str(), allocator), allocator);
    }
    {%-for topic_param in config.topics.params %}
    if ({{topic_param|lowerCamelCase}}.has_value())
    {
        parent.AddMember("{{topic_param}}", rapidjson::Value({{topic_param|lowerCamelCase}}.value().c_str(), allocator), allocator);
    }
    {%-endfor%}
}

{{stinger.name|UpperCamelCase}}Discovery::{{stinger.name|UpperCamelCase}}Discovery(std::shared_ptr<stinger::utils::IConnection> broker)
    : _broker(broker)
{
    // Subscribe to the discovery topic
    std::map<std::string, std::string> topicArgs;
    topicArgs["service_id"] = "+";
    topicArgs["interface_name"] = "{{stinger.name}}";
    topicArgs["client_id"] = "+";
    {%-for topic_param in config.topics.params %}
    topicArgs["{{topic_param}}"] = "+";
    {%-endfor%}

    _discoverySubscriptionId = _broker->Subscribe(stinger::utils::format("{{stinger.interface_info_topic()}}", topicArgs), 1);
    _allPropertySubscriptionId = _broker->Subscribe(stinger::utils::format("{{stinger.all_properties_value_topic()}}", topicArgs), 1);

    // Register message callback
    _brokerMessageCallbackHandle = _broker->AddMessageCallback([this](const stinger::mqtt::Message& msg) {
        _onMessage(msg);
    });
}

{{stinger.name|UpperCamelCase}}Discovery::~{{stinger.name|UpperCamelCase}}Discovery()
{
    if (_broker && _brokerMessageCallbackHandle != 0)
    {
        _broker->RemoveMessageCallback(_brokerMessageCallbackHandle);
        _brokerMessageCallbackHandle = 0;
    }
}

void {{stinger.name|UpperCamelCase}}Discovery::SetDiscoveryCallback(const std::function<void(const InstanceInfo&)>& cb)
{
    std::lock_guard<std::mutex> lock(_mutex);
    _discovery_callback = cb;
}

std::future<InstanceInfo> {{stinger.name|UpperCamelCase}}Discovery::GetSingleton()
{
    std::lock_guard<std::mutex> lock(_mutex);

    // If we already have at least one instance, return it immediately
    if (!_discoveredInstances.empty()) {
        std::promise<InstanceInfo> promise;
        promise.set_value(_discoveredInstances.begin()->second);
        return promise.get_future();
    }

    // Otherwise, create a promise that will be fulfilled when an instance is discovered
    _pending_promises.emplace_back();
    return _pending_promises.back().get_future();
}

std::vector<InstanceInfo> {{stinger.name|UpperCamelCase}}Discovery::GetInstances() const
{
    std::lock_guard<std::mutex> lock(_mutex);
    std::vector<InstanceInfo> instances;
    for (const auto& [key, info] : _discoveredInstances)
    {
        instances.push_back(info);
    }
    return instances;
}

void {{stinger.name|UpperCamelCase}}Discovery::_onMessage(const stinger::mqtt::Message& msg)
{
    // Check content type
    if (!msg.properties.contentType.has_value() || msg.properties.contentType.value() != "application/json")
    {
        std::cerr << "Invalid content type in Discovery message. Expected 'application/json'" << std::endl;
        return;
    }

    // Parse the JSON payload
    rapidjson::Document document;
    document.Parse(msg.payload.c_str());

    if (document.HasParseError())
    {
        std::cerr << "JSON parse error in Discovery: " 
                  << rapidjson::GetParseError_En(document.GetParseError())
                  << " at offset " << document.GetErrorOffset() << std::endl;
        return;
    }

    InstanceInfo *infoPtr = nullptr;

    if (msg.properties.subscriptionId == _discoverySubscriptionId)
    {
        std::vector<std::string> hashableIdentifiers;
        const uint8_t instance_id_expected_index = {{utils.topic_template_placeholder_index(config.topics.interface_discovery, "service_id")}};
        {%-for topic_param in config.topics.params %}
        const uint8_t {{topic_param}}_expected_index = {{utils.topic_template_placeholder_index(config.topics.interface_discovery, topic_param)}};
        {%endfor%}

        // Split topic by '/' into vector of strings
        std::vector<std::string> topicParts;
        std::stringstream ss(msg.topic);
        std::string part;
        uint8_t index = 0;
        while (std::getline(ss, part, '/'))
        {
            topicParts.push_back(part);
            if ((index == instance_id_expected_index){%-for topic_param in config.topics.params %} || (index == {{topic_param}}_expected_index){%endfor%})
            {
                hashableIdentifiers.push_back(part);
            }
            index++;
        }
        
        uint64_t instanceHash = stinger::utils::hashStrings(hashableIdentifiers);

        if (_discoveredInstances.find(instanceHash) == _discoveredInstances.end())
        {
            // Create a new InstanceInfo
            InstanceInfo info;
            info.UpdateFromRapidJsonObject(document);
            _discoveredInstances[instanceHash] = info;
            infoPtr = &_discoveredInstances[instanceHash];
        }
        else
        {
            // Update existing InstanceInfo with any new information from the payload
            _discoveredInstances[instanceHash].UpdateFromRapidJsonObject(document);
            infoPtr = &_discoveredInstances[instanceHash];
        }
    }{%if stinger.properties|length > 0%}
    else if (msg.properties.subscriptionId == _allPropertySubscriptionId)
    {
        // Parse out identifying information from the topic to find the corresponding InstanceInfo
        std::vector<std::string> hashableIdentifiers;
        std::string propertyName;
        uint32_t propertyVersion = 0;
        if (msg.properties.propertyVersion.has_value()) {
            propertyVersion = msg.properties.propertyVersion.value();
        }
        const uint8_t instance_id_expected_index = {{utils.topic_template_placeholder_index(config.topics.property_values, "service_id")}};
        {%-for topic_param in config.topics.params %}
        const uint8_t {{topic_param}}_expected_index = {{utils.topic_template_placeholder_index(config.topics.property_values, topic_param)}};
        {%endfor%}
        const uint8_t property_name_expected_index = {{utils.topic_template_placeholder_index(config.topics.property_values, "property_name")}};

        // Split topic by '/' into vector of strings
        std::vector<std::string> topicParts;
        std::stringstream ss(msg.topic);
        std::string part;
        uint8_t index = 0;
        while (std::getline(ss, part, '/'))
        {
            topicParts.push_back(part);
            if ((index == instance_id_expected_index){%-for topic_param in config.topics.params %} || (index == {{topic_param}}_expected_index){%endfor%})
            {
                hashableIdentifiers.push_back(part);
            }
            else if (index == property_name_expected_index)
            {
                propertyName = part;
            }
            index++;
        }
        
        uint64_t instanceHash = stinger::utils::hashStrings(hashableIdentifiers);
        
        auto it = _discoveredInstances.find(instanceHash);
        if (it == _discoveredInstances.end())
        {
            InstanceInfo info;
            _discoveredInstances[instanceHash] = info;
            infoPtr = &_discoveredInstances[instanceHash];
        }
        else
        {
            infoPtr = &it->second;
        }

        if (infoPtr)
        {
            {%for prop_name, prop in stinger.properties.items()%}
            {%if not loop.first%}else {%endif%}if (propertyName == "{{prop_name}}")
            {
                auto propertyObj = {{prop.cpp.property_struct_name}}::FromRapidJsonObject(document);
                infoPtr->initial_property_values.{{prop_name | lowerCamelCase}} = propertyObj;
                infoPtr->initial_property_values.{{prop_name | lowerCamelCase}}Version = propertyVersion;
            }
            {%endfor%}
        }

    }
    {%endif%}{# has any properties #}

    if (infoPtr && infoPtr->isComplete())
    {
        std::lock_guard<std::mutex> lock(_mutex);

        // Fulfill any pending promises
        for (auto& promise : _pending_promises)
        {
            promise.set_value(*infoPtr);
        }
        _pending_promises.clear();

        // Call the discovery callback if set
        if (_discovery_callback)
        {
            _discovery_callback(*infoPtr);
        }
    }
} // end of _onMessage

} // namespace {{stinger.name|snake_case}}
{%for ns in config.cpp.namespace|reverse%}
} // namespace {{ns}}
{%endfor-%}
