smartbotic/src/runner/node_registry.cpp

#include "node_registry.hpp"
#include "logging/logger.hpp"
#include "common/time_utils.hpp"

namespace smartbotic::runner {

using namespace common;

nlohmann::json NodeDefinition::toJson() const {
    nlohmann::json j;
    j["id"] = id;
    j["name"] = name;
    j["category"] = category;
    j["version"] = version;
    j["description"] = description;
    j["icon"] = icon;
    j["isTrigger"] = is_trigger;
    j["configSchema"] = config_schema;
    j["inputSchema"] = input_schema;
    j["outputSchema"] = output_schema;

    j["inputs"] = nlohmann::json::array();
    for (const auto& input : inputs) {
        j["inputs"].push_back({
            {"name", input.name},
            {"displayName", input.display_name},
            {"type", input.type},
            {"required", input.required}
        });
    }

    j["outputs"] = nlohmann::json::array();
    for (const auto& output : outputs) {
        nlohmann::json out_obj = {
            {"name", output.name},
            {"displayName", output.display_name},
            {"type", output.type}
        };
        if (!output.color.empty()) {
            out_obj["color"] = output.color;
        }
        j["outputs"].push_back(out_obj);
    }

    return j;
}

NodeDefinition NodeDefinition::fromProto(const proto::NodeDefinitionWithCode& proto) {
    NodeDefinition node;
    node.id = proto.id();
    node.name = proto.name();
    node.category = proto.category();
    node.version = proto.version();
    node.description = proto.description();
    node.icon = proto.icon();
    node.code = proto.code();
    node.is_trigger = proto.is_trigger();
    node.last_modified = proto.updated_at();

    // Parse schemas
    try {
        if (!proto.config_schema().empty() && proto.config_schema() != "null") {
            node.config_schema = nlohmann::json::parse(proto.config_schema());
        }
    } catch (...) {}

    try {
        if (!proto.input_schema().empty() && proto.input_schema() != "null") {
            node.input_schema = nlohmann::json::parse(proto.input_schema());
        }
    } catch (...) {}

    try {
        if (!proto.output_schema().empty() && proto.output_schema() != "null") {
            node.output_schema = nlohmann::json::parse(proto.output_schema());
        }
    } catch (...) {}

    // Parse inputs
    for (const auto& input : proto.inputs()) {
        NodeIO io;
        io.name = input.name();
        io.display_name = input.display_name();
        io.type = input.type();
        io.required = input.required();
        node.inputs.push_back(io);
    }

    // Parse outputs
    for (const auto& output : proto.outputs()) {
        NodeIO io;
        io.name = output.name();
        io.display_name = output.display_name();
        io.type = output.type();
        io.color = output.color();
        node.outputs.push_back(io);
    }

    return node;
}

NodeRegistry::NodeRegistry(const NodeRegistryConfig& config)
    : config_(config) {
    // Create gRPC channel
    channel_ = grpc::CreateChannel(config_.webserver_address, grpc::InsecureChannelCredentials());
    stub_ = proto::NodeSyncService::NewStub(channel_);
}

NodeRegistry::~NodeRegistry() {
    stop();
}

void NodeRegistry::start() {
    if (running_ || !config_.sync_enabled) {
        return;
    }

    running_ = true;

    // Load initial nodes
    auto result = loadNodesFromWebServer();
    if (result.failed()) {
        LOG_WARN("Failed to load initial nodes: {}", result.error().message());
    }

    // Start subscription thread for live updates
    subscribe_thread_ = std::thread(&NodeRegistry::subscribeLoop, this);
    LOG_INFO("Node registry sync started with {}", config_.webserver_address);
}

void NodeRegistry::stop() {
    if (!running_) {
        return;
    }

    LOG_INFO("Stopping node registry...");
    running_ = false;

    // Cancel any active gRPC stream to unblock subscribeLoop
    {
        std::lock_guard<std::mutex> lock(context_mutex_);
        if (active_context_) {
            active_context_->TryCancel();
        }
    }

    if (subscribe_thread_.joinable()) {
        subscribe_thread_.join();
    }

    LOG_INFO("Node registry sync stopped");
}

Result<void> NodeRegistry::loadNodesFromWebServer() {
    proto::GetAllNodesRequest request;
    proto::GetAllNodesResponse response;

    grpc::ClientContext context;
    context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(30));

    auto status = stub_->GetAllNodes(&context, request, &response);

    if (!status.ok()) {
        connected_ = false;
        return Error(ErrorCode::Internal, "Failed to get nodes: " + status.error_message());
    }

    connected_ = true;

    std::unique_lock lock(mutex_);
    nodes_.clear();

    for (const auto& proto_node : response.nodes()) {
        NodeDefinition node = NodeDefinition::fromProto(proto_node);
        nodes_[node.id] = node;
        LOG_DEBUG("Loaded node: {} ({})", node.id, node.name);
    }

    LOG_INFO("Loaded {} node definitions from webserver", nodes_.size());
    return {};
}

void NodeRegistry::subscribeLoop() {
    while (running_) {
        proto::SubscribeToNodeChangesRequest request;
        request.set_runner_id("runner_" + std::to_string(reinterpret_cast<uintptr_t>(this)));

        grpc::ClientContext context;

        // Store context pointer so stop() can cancel it
        {
            std::lock_guard<std::mutex> lock(context_mutex_);
            active_context_ = &context;
        }

        auto reader = stub_->SubscribeToNodeChanges(&context, request);

        if (!reader) {
            {
                std::lock_guard<std::mutex> lock(context_mutex_);
                active_context_ = nullptr;
            }
            if (!running_) break;
            LOG_WARN("Failed to subscribe to node changes, retrying...");
            std::this_thread::sleep_for(std::chrono::milliseconds(config_.reconnect_interval_ms));
            continue;
        }

        connected_ = true;
        LOG_INFO("Subscribed to node changes from webserver");

        proto::NodeChangeEvent event;
        while (running_ && reader->Read(&event)) {
            handleNodeChange(event);
        }

        // Clear context pointer
        {
            std::lock_guard<std::mutex> lock(context_mutex_);
            active_context_ = nullptr;
        }

        connected_ = false;
        auto status = reader->Finish();

        if (!status.ok() && running_) {
            LOG_WARN("Node change subscription ended: {}, reconnecting...", status.error_message());
            std::this_thread::sleep_for(std::chrono::milliseconds(config_.reconnect_interval_ms));
        }
    }
}

void NodeRegistry::handleNodeChange(const proto::NodeChangeEvent& event) {
    switch (event.type()) {
        case proto::NodeChangeEvent::CHANGE_TYPE_CREATED:
        case proto::NodeChangeEvent::CHANGE_TYPE_UPDATED: {
            NodeDefinition node = NodeDefinition::fromProto(event.node());
            {
                std::unique_lock lock(mutex_);
                nodes_[node.id] = node;
            }
            LOG_INFO("Node {}: {}",
                     event.type() == proto::NodeChangeEvent::CHANGE_TYPE_CREATED ? "created" : "updated",
                     node.id);

            if (reload_callback_) {
                reload_callback_(node.id);
            }
            break;
        }

        case proto::NodeChangeEvent::CHANGE_TYPE_DELETED: {
            {
                std::unique_lock lock(mutex_);
                nodes_.erase(event.node_id());
            }
            LOG_INFO("Node deleted: {}", event.node_id());

            if (reload_callback_) {
                reload_callback_(event.node_id());
            }
            break;
        }

        default:
            break;
    }
}

std::optional<NodeDefinition> NodeRegistry::getNode(const std::string& id) const {
    std::shared_lock lock(mutex_);

    auto it = nodes_.find(id);
    if (it != nodes_.end()) {
        return it->second;
    }
    return std::nullopt;
}

std::vector<NodeDefinition> NodeRegistry::getAllNodes() const {
    std::shared_lock lock(mutex_);

    std::vector<NodeDefinition> result;
    result.reserve(nodes_.size());

    for (const auto& [id, node] : nodes_) {
        result.push_back(node);
    }

    return result;
}

std::vector<NodeDefinition> NodeRegistry::getNodesByCategory(const std::string& category) const {
    std::shared_lock lock(mutex_);

    std::vector<NodeDefinition> result;

    for (const auto& [id, node] : nodes_) {
        if (node.category == category) {
            result.push_back(node);
        }
    }

    return result;
}

std::vector<NodeDefinition> NodeRegistry::getTriggerNodes() const {
    std::shared_lock lock(mutex_);

    std::vector<NodeDefinition> result;

    for (const auto& [id, node] : nodes_) {
        if (node.is_trigger) {
            result.push_back(node);
        }
    }

    return result;
}

Result<std::string> NodeRegistry::getNodeCode(const std::string& id) const {
    std::shared_lock lock(mutex_);

    auto it = nodes_.find(id);
    if (it == nodes_.end()) {
        return Error(ErrorCode::NotFound, "Node not found: " + id);
    }

    return it->second.code;
}

void NodeRegistry::setReloadCallback(NodeReloadCallback callback) {
    reload_callback_ = std::move(callback);
}

} // namespace smartbotic::runner