daggy/daggy/src/Serialization.cpp

#include <sstream>
#include <iomanip>

#include <daggy/Serialization.hpp>
#include <daggy/Utilities.hpp>

namespace daggy {

    ParameterValues parametersFromJSON(const std::string &jsonSpec) {
        rj::Document doc;
        rj::ParseResult parseResult = doc.Parse(jsonSpec.c_str());
        if (!parseResult) {
            throw std::runtime_error("Parameters spec is not valid JSON");
        }
        return parametersFromJSON(doc);
    }

    ParameterValues parametersFromJSON(const rj::Document &spec) {
        std::unordered_map<std::string, ParameterValue> parameters;
        if (!spec.IsObject()) { throw std::runtime_error("Parameters in spec is not a JSON dictionary"); }
        for (auto it = spec.MemberBegin(); it != spec.MemberEnd(); ++it) {
            if (!it->name.IsString()) {
                throw std::runtime_error("All keys must be strings.");
            }
            std::string name = std::string{"{{"} + it->name.GetString() + "}}";
            if (it->value.IsArray()) {
                std::vector<std::string> values;
                for (size_t i = 0; i < it->value.Size(); ++i) {
                    if (!it->value[i].IsString()) {
                        throw std::runtime_error(
                                "Attribute for " + std::string{it->name.GetString()} + " item " + std::to_string(i) +
                                " is not a string.");
                    }
                    values.emplace_back(it->value[i].GetString());
                }
                parameters[name] = values;
            } else if (it->value.IsString()) {
                parameters[name] = it->value.GetString();
            } else {
                throw std::runtime_error(
                        "Attribute for " + std::string{it->name.GetString()} + " is not a string or an array.");
            }
        }
        return parameters;
    }

    std::vector<Task> tasksFromJSON(const std::string &jsonSpec, const ParameterValues &parameters) {
        rj::Document doc;
        rj::ParseResult parseResult = doc.Parse(jsonSpec.c_str());
        if (!parseResult) {
            throw std::runtime_error("Unable to parse spec: ");
        }
        return tasksFromJSON(doc, parameters);
    }

    std::vector<Task> tasksFromJSON(const rj::Document &spec, const ParameterValues &parameters) {
        std::vector<Task> tasks;
        if (!spec.IsArray()) { throw std::runtime_error("Tasks is not an array"); }

        const std::vector<std::string> reqFields{"name", "command"};
        std::unordered_map<std::string, std::vector<std::string>> childrenMap;
        // Maps child -> parent
        std::unordered_map<std::string, std::vector<std::string>> parentMap;
        std::unordered_map<std::string, size_t> taskIndex;

        // Tasks
        for (size_t i = 0; i < spec.Size(); ++i) {
            if (!spec[i].IsObject()) {
                throw std::runtime_error("Task " + std::to_string(i) + " is not a dictionary.");
            }
            const auto &taskSpec = spec[i].GetObject();

            for (const auto &reqField : reqFields) {
                if (!taskSpec.HasMember(reqField.c_str())) {
                    throw std::runtime_error("Task " + std::to_string(i) + " is missing required field " + reqField);
                }
            }

            // Grab the standard fields with defaults;
            std::string name = taskSpec["name"].GetString();
            taskIndex[name] = i;

            uint8_t maxRetries = 0;
            if (taskSpec.HasMember("maxRetries")) { maxRetries = taskSpec["maxRetries"].GetInt(); }
            uint8_t retryIntervalSeconds = 0;
            if (taskSpec.HasMember(
                    "retryIntervalSeconds")) { retryIntervalSeconds = taskSpec["retryIntervalSeconds"].GetInt(); }

            // Children / parents
            std::unordered_set<std::string> children;
            if (taskSpec.HasMember("children")) {
                const auto &specChildren = taskSpec["children"].GetArray();
                for (size_t c = 0; c < specChildren.Size(); ++c) {
                    children.insert(specChildren[c].GetString());
                }
            }
            if (taskSpec.HasMember("parents")) {
                const auto &specParents = taskSpec["parents"].GetArray();
                for (size_t c = 0; c < specParents.Size(); ++c) {
                    parentMap[name].emplace_back(specParents[c].GetString());
                }
            }

            // Build out the commands
            std::vector<std::string> command;
            for (size_t cmd = 0; cmd < taskSpec["command"].Size(); ++cmd) {
                command.emplace_back(taskSpec["command"][cmd].GetString());
            }
            auto commands = expandCommands(command, parameters);

            // Create the tasks
            auto &taskNames = childrenMap[name];
            for (size_t tid = 0; tid < commands.size(); ++tid) {
                std::string taskName = name + "_" + std::to_string(tid);
                taskNames.push_back(taskName);
                tasks.emplace_back(Task{
                        .name = name + "_" + std::to_string(tid),
                        .command = commands[tid],
                        .maxRetries = maxRetries,
                        .retryIntervalSeconds = retryIntervalSeconds,
                        .children = children
                });
            }
        }

        // Update any missing child -> parent relationship
        for (auto &task : tasks) {
            auto pit = parentMap.find(task.name);
            if (pit == parentMap.end()) { continue; }

            for (const auto &parent : pit->second) {
                tasks[taskIndex[parent]].children.insert(task.name);
            }
        }

        // At the end, replace the names of the children with all the expanded versions
        for (auto &task : tasks) {
            std::unordered_set<std::string> children;
            for (const auto &child : task.children) {
                auto &newChildren = childrenMap[child];
                std::copy(newChildren.begin(), newChildren.end(), std::inserter(children, children.end()));
            }
            task.children.swap(children);
        }

        return tasks;
    }

    // I really want to do this with rapidjson, but damn they make it ugly and difficult.
    // So we'll shortcut and generate the JSON directly.
    std::string taskToJSON(const Task &task) {
        std::stringstream ss;
        bool first = false;

        ss << "{"
           << R"("name": )" << std::quoted(task.name) << ','
           << R"("maxRetries": )" << task.maxRetries << ','
           << R"("retryIntervalSeconds": )" << task.retryIntervalSeconds << ',';

        // Commands
        ss << R"("command": [)";
        first = true;
        for (const auto &part : task.command) {
            if (!first) ss << ',';
            ss << std::quoted(part);
            first = false;
        }
        ss << "],";

        ss << R"("children": [)";
        first = true;
        for (const auto &child : task.children) {
            if (!first) ss << ',';
            ss << std::quoted(child);
            first = false;
        }
        ss << "]";

        ss << '}';
        return ss.str();
    }

    std::string tasksToJSON(const std::vector<Task> &tasks) {
        std::stringstream ss;

        ss << "[";

        bool first = true;
        for (const auto &task : tasks) {
            if (!first) ss << ',';
            ss << taskToJSON(task);
            first = false;
        }
        ss << "]";

        return ss.str();
    }
}