smartbotic/src/webserver/websocket_server.cpp

#include "websocket_server.hpp"
#include "common/uuid.hpp"
#include "logging/logger.hpp"
#include <cstring>

namespace smartbotic::webserver {

// Per-session data
struct PerSessionData {
    WebSocketClient* client;
};

// Forward declarations for LWS callbacks
static int ws_callback(struct lws* wsi, enum lws_callback_reasons reason,
                       void* user, void* in, size_t len);

// Global pointer for callback access (set per-context in user data)
static WebSocketServer* g_ws_server = nullptr;

// Protocol definition
static struct lws_protocols protocols[] = {
    {
        "smartbotic",
        ws_callback,
        sizeof(PerSessionData),
        1024 * 64,  // rx buffer size
        0, nullptr, 0
    },
    { nullptr, nullptr, 0, 0, 0, nullptr, 0 }
};

static int ws_callback(struct lws* wsi, enum lws_callback_reasons reason,
                       void* user, void* in, size_t len) {
    auto* pss = static_cast<PerSessionData*>(user);

    switch (reason) {
        case LWS_CALLBACK_ESTABLISHED:
            if (g_ws_server) {
                g_ws_server->onConnect(wsi);
            }
            break;

        case LWS_CALLBACK_CLOSED:
            if (g_ws_server) {
                g_ws_server->onDisconnect(wsi);
            }
            break;

        case LWS_CALLBACK_RECEIVE:
            if (g_ws_server && in && len > 0) {
                return g_ws_server->onReceive(wsi, static_cast<const char*>(in), len);
            }
            break;

        case LWS_CALLBACK_SERVER_WRITEABLE:
            if (g_ws_server) {
                return g_ws_server->onWritable(wsi);
            }
            break;

        default:
            break;
    }

    return 0;
}

WebSocketServer::WebSocketServer(const WebSocketServerConfig& config, auth::JwtUtils& jwt)
    : config_(config), jwt_(jwt) {
    g_ws_server = this;
}

WebSocketServer::~WebSocketServer() {
    stop();
    g_ws_server = nullptr;
}

void WebSocketServer::start() {
    if (running_) {
        return;
    }

    struct lws_context_creation_info info{};
    memset(&info, 0, sizeof(info));

    info.port = config_.port;
    info.protocols = protocols;
    info.gid = -1;
    info.uid = -1;
    info.options = LWS_SERVER_OPTION_HTTP_HEADERS_SECURITY_BEST_PRACTICES_ENFORCE;

    context_ = lws_create_context(&info);
    if (!context_) {
        LOG_ERROR("Failed to create WebSocket context");
        return;
    }

    running_ = true;
    service_thread_ = std::thread(&WebSocketServer::serviceLoop, this);

    LOG_INFO("WebSocket server started on port {}", config_.port);
}

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

    running_ = false;

    // Cancel the service loop immediately
    if (context_) {
        lws_cancel_service(context_);
    }

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

    if (context_) {
        lws_context_destroy(context_);
        context_ = nullptr;
    }

    LOG_INFO("WebSocket server stopped");
}

void WebSocketServer::serviceLoop() {
    while (running_) {
        lws_service(context_, 50);  // 50ms timeout

        // Process any pending writable requests (thread-safe wakeup mechanism)
        std::vector<struct lws*> to_signal;
        {
            std::lock_guard<std::mutex> lock(pending_mutex_);
            to_signal.assign(pending_writable_.begin(), pending_writable_.end());
            pending_writable_.clear();
        }

        for (auto* wsi : to_signal) {
            lws_callback_on_writable(wsi);
        }
    }
}

int WebSocketServer::onConnect(struct lws* wsi) {
    std::unique_lock lock(clients_mutex_);

    auto client = std::make_unique<WebSocketClient>();
    client->wsi = wsi;
    client->id = common::UUID::generate();

    client_id_map_[client->id] = wsi;
    clients_[wsi] = std::move(client);

    LOG_DEBUG("WebSocket client connected: {}", clients_[wsi]->id);
    return 0;
}

void WebSocketServer::onDisconnect(struct lws* wsi) {
    // Which workflows this client was watching has to be read before it is
    // erased, but the roster can only be republished after - and publishing
    // takes the same lock. So: collect, drop the lock, then publish.
    std::vector<std::string> was_watching;
    std::string gone_client_id;

    {
        std::unique_lock lock(clients_mutex_);

        auto it = clients_.find(wsi);
        if (it != clients_.end()) {
            LOG_DEBUG("WebSocket client disconnected: {}", it->second->id);
            gone_client_id = it->second->id;
            for (const auto& sub : it->second->subscriptions) {
                std::string workflow_id = presenceWorkflowId(sub);
                if (!workflow_id.empty()) was_watching.push_back(workflow_id);
            }
            client_id_map_.erase(it->second->id);
            clients_.erase(it);
        }
    }

    // A lock outlives its holder only if nothing notices they have gone, which
    // is exactly how an editor ends up with a node nobody can touch again.
    std::vector<std::string> unlocked;
    releaseLocksOf(gone_client_id, &unlocked);

    for (const auto& workflow_id : was_watching) {
        publishPresence(workflow_id);
    }
    for (const auto& workflow_id : unlocked) {
        publishLocks(workflow_id);
    }
}

int WebSocketServer::onReceive(struct lws* wsi, const char* data, size_t len) {
    WebSocketClient* client = nullptr;

    // Get client pointer under lock, then release before processing
    {
        std::shared_lock lock(clients_mutex_);
        auto it = clients_.find(wsi);
        if (it == clients_.end()) {
            return 0;
        }
        client = it->second.get();
    }

    // Process message without holding the lock (processMessage may call sendToClient
    // which needs unique_lock, so we must not hold shared_lock here)
    try {
        std::string str(data, len);
        auto message = nlohmann::json::parse(str);
        processMessage(*client, message);
    } catch (const std::exception& e) {
        LOG_WARN("Failed to parse WebSocket message: {}", e.what());
    }

    return 0;
}

int WebSocketServer::onWritable(struct lws* wsi) {
    std::unique_lock lock(clients_mutex_);

    auto it = clients_.find(wsi);
    if (it == clients_.end() || it->second->send_queue.empty()) {
        return 0;
    }

    auto& queue = it->second->send_queue;
    std::string& msg = queue.front();

    std::vector<unsigned char> buf(LWS_PRE + msg.size());
    memcpy(buf.data() + LWS_PRE, msg.data(), msg.size());

    int written = lws_write(wsi, buf.data() + LWS_PRE, msg.size(), LWS_WRITE_TEXT);
    if (written < 0) {
        LOG_ERROR("WebSocket write failed");
        return -1;
    }

    queue.erase(queue.begin());

    if (!queue.empty()) {
        lws_callback_on_writable(wsi);
    }

    return 0;
}

void WebSocketServer::processMessage(WebSocketClient& client, const nlohmann::json& message) {
    std::string type = message.value("type", "");

    if (type == "auth") {
        handleAuth(client, message);
    } else if (type == "subscribe") {
        handleSubscribe(client, message);
    } else if (type == "unsubscribe") {
        handleUnsubscribe(client, message);
    } else if (type == "lock") {
        handleLock(client, message);
    } else if (type == "unlock") {
        handleUnlock(client, message);
    } else if (type == "node_moved") {
        handleNodeMoved(client, message);
    } else if (type == "graph_edit") {
        handleGraphEdit(client, message);
    } else if (message_handler_) {
        message_handler_(client.id, message);
    }
}

void WebSocketServer::handleAuth(WebSocketClient& client, const nlohmann::json& message) {
    std::string token = message.value("token", "");

    // Remove "Bearer " prefix if present
    if (token.starts_with("Bearer ")) {
        token = token.substr(7);
    }

    auto result = jwt_.verifyToken(token);
    if (result.ok()) {
        client.authenticated = true;
        client.user_id = result.value().user_id;
        client.username = result.value().username;

        nlohmann::json response;
        response["type"] = "auth_success";
        response["userId"] = client.user_id;
        response["username"] = client.username;
        // This connection's own id. A lock is per connection, so an editor has
        // to be able to tell its own lock from one held by the same person in
        // another tab.
        response["clientId"] = client.id;
        sendToClient(client.id, response);

        LOG_DEBUG("WebSocket client authenticated: {}", client.id);
    } else {
        nlohmann::json response;
        response["type"] = "auth_error";
        response["error"] = result.error().message();
        sendToClient(client.id, response);
    }
}

void WebSocketServer::handleSubscribe(WebSocketClient& client, const nlohmann::json& message) {
    if (!client.authenticated) {
        nlohmann::json response;
        response["type"] = "error";
        response["error"] = "Not authenticated";
        sendToClient(client.id, response);
        return;
    }

    auto channels = message.value("channels", std::vector<std::string>{});
    for (const auto& channel : channels) {
        subscribe(client.id, channel);
        // Publishing after subscribing means the arriving client is itself on
        // the channel, so it gets the roster as its first message rather than
        // waiting for somebody else to come or go.
        std::string workflow_id = presenceWorkflowId(channel);
        if (!workflow_id.empty()) publishPresence(workflow_id);

        // The same for locks, which are otherwise only published when they
        // change: an editor arriving after somebody started working would see
        // an unlocked canvas and walk straight into a node already taken.
        std::string locks_workflow = channelWorkflowId(channel, ".locks");
        if (!locks_workflow.empty()) publishLocks(locks_workflow);
    }
}

void WebSocketServer::handleUnsubscribe(WebSocketClient& client, const nlohmann::json& message) {
    auto channels = message.value("channels", std::vector<std::string>{});
    for (const auto& channel : channels) {
        unsubscribe(client.id, channel);
        std::string workflow_id = presenceWorkflowId(channel);
        if (!workflow_id.empty()) publishPresence(workflow_id);
    }
}

std::string WebSocketServer::presenceChannel(const std::string& workflow_id) {
    return "workflows." + workflow_id + ".presence";
}

std::string WebSocketServer::channelWorkflowId(const std::string& channel,
                                              const std::string& suffix) {
    static const std::string prefix = "workflows.";
    if (!channel.starts_with(prefix) || !channel.ends_with(suffix)) return "";

    // A wildcard subscription is a listener, not somebody with the workflow
    // open, and counting it would put phantom names in the roster.
    std::string id = channel.substr(prefix.size(),
                                    channel.size() - prefix.size() - suffix.size());
    if (id.empty() || id.find('*') != std::string::npos) return "";
    return id;
}

std::string WebSocketServer::presenceWorkflowId(const std::string& channel) {
    return channelWorkflowId(channel, ".presence");
}

void WebSocketServer::publishPresence(const std::string& workflow_id) {
    const std::string channel = presenceChannel(workflow_id);

    nlohmann::json viewers = nlohmann::json::array();
    {
        std::shared_lock lock(clients_mutex_);

        // One person with the editor open in two tabs is one person, so the
        // roster is by user rather than by connection.
        std::unordered_map<std::string, size_t> seen;
        for (const auto& [wsi, client] : clients_) {
            (void)wsi;
            if (!client->authenticated) continue;
            if (!client->subscriptions.contains(channel)) continue;

            auto it = seen.find(client->user_id);
            if (it != seen.end()) {
                viewers[it->second]["connections"] =
                    viewers[it->second]["connections"].get<int>() + 1;
                continue;
            }
            seen.emplace(client->user_id, viewers.size());
            viewers.push_back({
                {"userId", client->user_id},
                {"username", client->username},
                {"connections", 1},
            });
        }
    }

    broadcast(channel, {{"workflowId", workflow_id}, {"viewers", viewers}});
}


// ---------------------------------------------------------------------------
// Node locks and live movement
//
// While somebody is dragging or configuring a node, everyone else is kept off
// that one node - not off the workflow. Two people working on different parts
// of the same flow is the normal case and should stay possible; two people
// dragging the same node is the one that produces nonsense.
// ---------------------------------------------------------------------------

void WebSocketServer::handleLock(WebSocketClient& client, const nlohmann::json& message) {
    if (!client.authenticated) return;

    const std::string workflow_id = message.value("workflowId", "");
    const std::string node_id = message.value("nodeId", "");
    const std::string kind = message.value("kind", "editing");
    if (workflow_id.empty() || node_id.empty()) return;

    bool granted = false;
    std::string holder_username;
    {
        std::lock_guard<std::mutex> lock(locks_mutex_);
        auto& nodes = locks_[workflow_id];
        auto it = nodes.find(node_id);
        if (it == nodes.end() || it->second.client_id == client.id) {
            // Re-claiming your own lock is how the kind changes from dragging
            // to editing without a release in between.
            nodes[node_id] = NodeLock{client.id, client.user_id, client.username, kind};
            granted = true;
        } else {
            holder_username = it->second.username;
        }
    }

    if (!granted) {
        // Said out loud rather than ignored: a click that silently does nothing
        // reads as the editor being broken.
        nlohmann::json response;
        response["type"] = "lock_denied";
        response["workflowId"] = workflow_id;
        response["nodeId"] = node_id;
        response["heldBy"] = holder_username;
        sendToClient(client.id, response);
        return;
    }

    publishLocks(workflow_id);
}

void WebSocketServer::handleUnlock(WebSocketClient& client, const nlohmann::json& message) {
    const std::string workflow_id = message.value("workflowId", "");
    const std::string node_id = message.value("nodeId", "");
    if (workflow_id.empty() || node_id.empty()) return;

    bool changed = false;
    {
        std::lock_guard<std::mutex> lock(locks_mutex_);
        auto wf = locks_.find(workflow_id);
        if (wf != locks_.end()) {
            auto it = wf->second.find(node_id);
            // Only the holder can let go - otherwise a stale release from an
            // editor that has already moved on would free somebody else's node.
            if (it != wf->second.end() && it->second.client_id == client.id) {
                wf->second.erase(it);
                changed = true;
            }
        }
    }
    if (changed) publishLocks(workflow_id);
}

void WebSocketServer::handleNodeMoved(WebSocketClient& client, const nlohmann::json& message) {
    if (!client.authenticated) return;

    const std::string workflow_id = message.value("workflowId", "");
    const std::string node_id = message.value("nodeId", "");
    if (workflow_id.empty() || node_id.empty()) return;

    // Only the editor holding the node may say where it is. Without this,
    // anyone could move anybody's node by sending the message directly.
    {
        std::lock_guard<std::mutex> lock(locks_mutex_);
        auto wf = locks_.find(workflow_id);
        if (wf == locks_.end()) return;
        auto it = wf->second.find(node_id);
        if (it == wf->second.end() || it->second.client_id != client.id) return;
    }

    // Not stored: a position in flight is worth nothing once it has been
    // delivered, and the authoritative one is whatever gets saved.
    broadcast("workflows." + workflow_id + ".moves", {
        {"workflowId", workflow_id},
        {"nodeId", node_id},
        {"position", message.value("position", nlohmann::json::object())},
        {"userId", client.user_id},
        {"clientId", client.id},
    });
}


void WebSocketServer::handleGraphEdit(WebSocketClient& client, const nlohmann::json& message) {
    if (!client.authenticated) return;

    const std::string workflow_id = message.value("workflowId", "");
    if (workflow_id.empty()) return;

    // Having the workflow open is the check that matters here. Anyone who can
    // see it can already save whatever they like to it, so demanding a lock per
    // edit would buy nothing and would make an auto-layout - which moves every
    // node at once and holds no lock at all - impossible to share.
    const std::string channel = "workflows." + workflow_id + ".edits";
    {
        std::shared_lock lock(clients_mutex_);
        if (!client.subscriptions.contains(channel)) return;
    }

    nlohmann::json out = message.value("delta", nlohmann::json::object());
    out["workflowId"] = workflow_id;
    // Stamped here rather than trusted from the sender, so an edit cannot be
    // attributed to somebody else.
    out["userId"] = client.user_id;
    out["username"] = client.username;
    out["clientId"] = client.id;

    broadcast(channel, out);
}

void WebSocketServer::releaseLocksOf(const std::string& client_id,
                                     std::vector<std::string>* touched_workflows) {
    std::lock_guard<std::mutex> lock(locks_mutex_);
    for (auto& [workflow_id, nodes] : locks_) {
        bool changed = false;
        for (auto it = nodes.begin(); it != nodes.end();) {
            if (it->second.client_id == client_id) {
                it = nodes.erase(it);
                changed = true;
            } else {
                ++it;
            }
        }
        if (changed && touched_workflows) touched_workflows->push_back(workflow_id);
    }
}

void WebSocketServer::publishLocks(const std::string& workflow_id) {
    nlohmann::json held = nlohmann::json::array();
    {
        std::lock_guard<std::mutex> lock(locks_mutex_);
        auto wf = locks_.find(workflow_id);
        if (wf != locks_.end()) {
            for (const auto& [node_id, holder] : wf->second) {
                held.push_back({
                    {"nodeId", node_id},
                    {"userId", holder.user_id},
                    {"username", holder.username},
                    {"clientId", holder.client_id},
                    {"kind", holder.kind},
                });
            }
        }
    }
    broadcast("workflows." + workflow_id + ".locks",
              {{"workflowId", workflow_id}, {"locks", held}});
}

void WebSocketServer::broadcast(const std::string& channel, const nlohmann::json& data) {
    nlohmann::json message;
    message["channel"] = channel;
    message["data"] = data;
    std::string msg_str = message.dump();

    std::vector<struct lws*> clients_to_signal;

    {
        std::unique_lock lock(clients_mutex_);  // Must be unique_lock - we modify send_queue

        for (auto& [wsi, client] : clients_) {
            if (!client->authenticated) {
                continue;
            }

            for (const auto& sub : client->subscriptions) {
                if (matchesChannel(sub, channel)) {
                    client->send_queue.push_back(msg_str);
                    clients_to_signal.push_back(wsi);
                    break;
                }
            }
        }
    }

    // Signal pending writes outside of clients_mutex_ to avoid deadlock
    if (!clients_to_signal.empty()) {
        std::lock_guard<std::mutex> lock(pending_mutex_);
        for (auto* wsi : clients_to_signal) {
            pending_writable_.insert(wsi);
        }
        // Wake up the service loop
        lws_cancel_service(context_);
    }

    LOG_DEBUG("Broadcast to channel {}: {} clients matched", channel, clients_to_signal.size());
}

void WebSocketServer::sendToClient(const std::string& client_id, const nlohmann::json& message) {
    struct lws* wsi_to_signal = nullptr;

    {
        std::unique_lock lock(clients_mutex_);

        auto it = client_id_map_.find(client_id);
        if (it == client_id_map_.end()) {
            return;
        }

        auto client_it = clients_.find(it->second);
        if (client_it == clients_.end()) {
            return;
        }

        client_it->second->send_queue.push_back(message.dump());
        wsi_to_signal = it->second;
    }

    // Signal pending write outside of clients_mutex_
    if (wsi_to_signal) {
        std::lock_guard<std::mutex> lock(pending_mutex_);
        pending_writable_.insert(wsi_to_signal);
        lws_cancel_service(context_);
    }
}

void WebSocketServer::sendToUser(const std::string& user_id, const nlohmann::json& message) {
    std::vector<struct lws*> clients_to_signal;
    std::string msg_str = message.dump();

    {
        std::unique_lock lock(clients_mutex_);

        for (auto& [wsi, client] : clients_) {
            if (client->user_id == user_id) {
                client->send_queue.push_back(msg_str);
                clients_to_signal.push_back(wsi);
            }
        }
    }

    // Signal pending writes outside of clients_mutex_
    if (!clients_to_signal.empty()) {
        std::lock_guard<std::mutex> lock(pending_mutex_);
        for (auto* wsi : clients_to_signal) {
            pending_writable_.insert(wsi);
        }
        lws_cancel_service(context_);
    }
}

void WebSocketServer::subscribe(const std::string& client_id, const std::string& channel) {
    std::unique_lock lock(clients_mutex_);

    auto it = client_id_map_.find(client_id);
    if (it == client_id_map_.end()) {
        return;
    }

    auto client_it = clients_.find(it->second);
    if (client_it == clients_.end()) {
        return;
    }

    client_it->second->subscriptions.insert(channel);
    LOG_DEBUG("Client {} subscribed to {}", client_id, channel);
}

void WebSocketServer::unsubscribe(const std::string& client_id, const std::string& channel) {
    std::unique_lock lock(clients_mutex_);

    auto it = client_id_map_.find(client_id);
    if (it == client_id_map_.end()) {
        return;
    }

    auto client_it = clients_.find(it->second);
    if (client_it == clients_.end()) {
        return;
    }

    client_it->second->subscriptions.erase(channel);
}

bool WebSocketServer::matchesChannel(const std::string& subscription, const std::string& channel) {
    // Support wildcard matching: "executions.*" matches "executions.123.started"
    if (subscription == channel) {
        return true;
    }

    if (subscription.ends_with(".*")) {
        std::string prefix = subscription.substr(0, subscription.length() - 1);
        return channel.starts_with(prefix);
    }

    return false;
}

void WebSocketServer::setMessageHandler(WebSocketMessageHandler handler) {
    message_handler_ = std::move(handler);
}

size_t WebSocketServer::getConnectionCount() const {
    std::shared_lock lock(clients_mutex_);
    return clients_.size();
}

} // namespace smartbotic::webserver