DMF-Studio/libs/dmf-node/src/control_server.cpp

#include <dmf-node/control_server.hpp>
#include <dmf-node/node.hpp>

#include <libwebsockets.h>

#include <nlohmann/json.hpp>

#include <spdlog/spdlog.h>

#include <cstring>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

namespace dmf_node {

struct ControlServerData {
    StatusCallback on_event;
    std::unordered_map<std::string, CommandHandler> commands;
    std::mutex send_mutex;
    std::vector<std::string> send_queue;
    struct lws* client_wsi = nullptr;
};

struct ControlServer::Impl {
    uint16_t port;
    std::unique_ptr<ControlServerData> data;
    struct lws_context* context = nullptr;
};

static nlohmann::json dispatch_command(ControlServerData* data, const nlohmann::json& msg) {
    if (!msg.contains("cmd")) {
        spdlog::warn("Control: message missing 'cmd' field");
        return {{"error", "missing 'cmd' field"}};
    }
    auto cmd = msg["cmd"].get<std::string>();
    auto it = data->commands.find(cmd);
    if (it != data->commands.end()) {
        return it->second(msg);
    } else {
        spdlog::warn("Control: unknown command '{}'", cmd);
        return {{"error", "unknown command: " + cmd}};
    }
}

static int send_http_json(struct lws* wsi, const std::string& status, const std::string& body) {
    auto hdr = "HTTP/1.1 " + status + "\r\n"
               "Content-Type: application/json\r\n"
               "Content-Length: " + std::to_string(body.size()) + "\r\n"
               "Connection: close\r\n"
               "\r\n";
    std::vector<uint8_t> buf(LWS_PRE + hdr.size() + body.size());
    std::memcpy(buf.data() + LWS_PRE, hdr.data(), hdr.size());
    std::memcpy(buf.data() + LWS_PRE + hdr.size(), body.data(), body.size());
    lws_write(wsi, buf.data() + LWS_PRE, hdr.size() + body.size(), LWS_WRITE_HTTP);
    lws_http_transaction_completed(wsi);
    return -1;
}

struct PerSession {
    std::string http_body;
    bool is_ws = false;
    ControlServerData* data = nullptr;
};

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

    switch (reason) {
    case LWS_CALLBACK_HTTP: {
        new (ps) PerSession();
        auto* vhost = lws_get_vhost(wsi);
        ps->data = vhost ? static_cast<ControlServerData*>(lws_vhost_user(vhost)) : nullptr;

        char* uri_ptr = nullptr;
        int uri_len = 0;
        int method = lws_http_get_uri_and_method(wsi, &uri_ptr, &uri_len);
        std::string path(uri_ptr ? uri_ptr : "", uri_len > 0 ? uri_len : 0);

        if (path == "/cmd" && method == LWSHUMETH_POST) {
            int cl = lws_hdr_total_length(wsi, WSI_TOKEN_HTTP_CONTENT_LENGTH);
            if (cl > 0) {
                ps->http_body.reserve(cl);
            }
            return 0;
        }

        if (path == "/cmd" && (method == LWSHUMETH_GET || method == -1)) {
            return send_http_json(wsi, "405 Method Not Allowed", R"({"error":"POST only"})");
        }

        return 0;
    }
    case LWS_CALLBACK_HTTP_BODY: {
        ps->http_body.append(static_cast<char*>(in), len);
        break;
    }
    case LWS_CALLBACK_HTTP_BODY_COMPLETION: {
        if (!ps->data) {
            return send_http_json(wsi, "500 Error", R"({"error":"no data"})");
        }
        try {
            auto msg = nlohmann::json::parse(ps->http_body);
            auto result = dispatch_command(ps->data, msg);
            return send_http_json(wsi, "200 OK", result.dump());
        } catch (const nlohmann::json::parse_error& e) {
            return send_http_json(wsi, "400 Bad Request",
                                 nlohmann::json({{"error", e.what()}}).dump());
        }
    }
    case LWS_CALLBACK_ESTABLISHED: {
        auto* vhost = lws_get_vhost(wsi);
        ps->data = vhost ? static_cast<ControlServerData*>(lws_vhost_user(vhost)) : nullptr;
        ps->is_ws = true;
        if (ps->data) {
            ps->data->client_wsi = wsi;
        }
        spdlog::info("Control WS: client connected");
        break;
    }
    case LWS_CALLBACK_RECEIVE: {
        if (!ps->data) {
            break;
        }
        try {
            auto msg = nlohmann::json::parse(static_cast<char*>(in), static_cast<char*>(in) + len);
            dispatch_command(ps->data, msg);
        } catch (const nlohmann::json::parse_error& e) {
            spdlog::warn("Control WS: JSON parse error: {}", e.what());
        }
        break;
    }
    case LWS_CALLBACK_SERVER_WRITEABLE: {
        if (!ps->data) {
            break;
        }
        std::lock_guard<std::mutex> lock(ps->data->send_mutex);
        while (!ps->data->send_queue.empty()) {
            auto& msg = ps->data->send_queue.back();
            std::vector<uint8_t> buf(LWS_PRE + msg.size());
            std::memcpy(buf.data() + LWS_PRE, msg.data(), msg.size());
            lws_write(wsi, buf.data() + LWS_PRE, msg.size(), LWS_WRITE_TEXT);
            ps->data->send_queue.pop_back();
        }
        break;
    }
    case LWS_CALLBACK_CLOSED: {
        if (ps->data) {
            ps->data->client_wsi = nullptr;
        }
        spdlog::info("Control WS: client disconnected");
        break;
    }
    default:
        break;
    }
    return 0;
}

static const struct lws_protocols protocols[] = {
    {
        "http-only",
        callback_all,
        sizeof(PerSession),
        0,
    },
    {
        "dmf-control",
        callback_all,
        sizeof(PerSession),
        65536,
    },
    {nullptr, nullptr, 0, 0},
};

static const struct lws_http_mount mounts[] = {
    {
        .mount_next = &mounts[1],
        .mountpoint = "/cmd",
        .origin = "",
        .def = "",
        .protocol = "http-only",
        .cgienv = nullptr,
        .extra_mimetypes = nullptr,
        .interpret = nullptr,
        .cgi_timeout = 0,
        .cache_max_age = 0,
        .auth_mask = 0,
        .cache_reusable = 0,
        .cache_revalidate = 0,
        .cache_intermediaries = 0,
        .origin_protocol = LWSMPRO_CALLBACK,
        .mountpoint_len = 4,
    },
    {
        .mount_next = nullptr,
        .mountpoint = "/",
        .origin = "",
        .def = "",
        .protocol = "dmf-control",
        .cgienv = nullptr,
        .extra_mimetypes = nullptr,
        .interpret = nullptr,
        .cgi_timeout = 0,
        .cache_max_age = 0,
        .auth_mask = 0,
        .cache_reusable = 0,
        .cache_revalidate = 0,
        .cache_intermediaries = 0,
        .origin_protocol = LWSMPRO_CALLBACK,
        .mountpoint_len = 1,
    },
};

ControlServer::ControlServer(uint16_t port, StatusCallback on_event)
    : impl_(std::make_unique<Impl>()) {
    impl_->port = port;
    impl_->data = std::make_unique<ControlServerData>();
    impl_->data->on_event = std::move(on_event);

    struct lws_context_creation_info info;
    std::memset(&info, 0, sizeof(info));
    info.port = port;
    info.protocols = protocols;
    info.mounts = mounts;
    info.user = impl_->data.get();
    info.gid = -1;
    info.uid = -1;

    impl_->context = lws_create_context(&info);
    if (!impl_->context) {
        spdlog::error("Failed to create WS context on port {}", port);
        throw std::runtime_error("Failed to create WS context");
    }
    spdlog::info("Control WS: listening on port {}", port);
}

ControlServer::~ControlServer() {
    if (impl_->context) {
        lws_context_destroy(impl_->context);
    }
}

void ControlServer::register_command(const std::string& cmd, CommandHandler handler) {
    impl_->data->commands[cmd] = std::move(handler);
}

void ControlServer::send_event(const nlohmann::json& event) {
    auto data = event.dump();
    {
        std::lock_guard<std::mutex> lock(impl_->data->send_mutex);
        impl_->data->send_queue.push_back(data);
    }
    if (impl_->data->client_wsi) {
        lws_callback_on_writable(impl_->data->client_wsi);
    }
}

void ControlServer::poll(int timeout_ms) {
    lws_service(impl_->context, timeout_ms);
}

} // namespace dmf_node