// See the file "COPYING" in the main distribution directory for copyright.

// A zero-dependency cluster-layout.zeek generator for single and multi-host
// Zeek deployments. Workers do not get ports allocated. Supports the multi-host
// use-case via -C /etc/zeek/cluster and will assemble a cluster layout based
// on per-host <hostname>.zeek.conf files in the given directory.
//
// For the single host use-case, expects the individual process counts as
// individual arguments, or passing it a <PREFIX>/etc/zeek/zeek.conf file
// via the -C flag.
//
// The cluster-layout.zeek content is produced on stdout or written to the
// path given to the -o flag.
//
// Usage Examples:
//
//     $ zeek-cluster-layout-generator -L <loggers> -P <proxies> -W <workers> -M <manager>
//         -a <listen_address> -p <listen_port>
//         -m <metrics_port> -b <metrics_address>
//         -x <cluster_node_prefix>
//         [-o outfile]
//
//     The -W flag also accepts a named form like -W 'eth0:2 eth1:4' that
//     will include the eth0 and eth1 interface names in the cluster node name.
//     The number following the colon is the number of workers.
//
//     Passing -m 0 will use metrics port 0/tcp for each node. This is useful
//     to disable metrics exposition via Prometheus.
//
//     # Create layout from a single host zeek.conf file
//     $ zeek-cluster-layout-generator -C /etc/zeek/zeek.conf [-o outfile]
//
//     # Create cluster layout from a multi-host cluster configuration directory
//     $ zeek-cluster-layout-generator -C /etc/zeek/cluster [-o outfile]
//
#include <unistd.h>
#include <algorithm>
#include <cctype>
#include <cerrno>
#include <charconv>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <regex>
#include <string>
#include <system_error>
#include <vector>

#include "zeek-cluster-config.h"

namespace {

int checked_stoi(int opt, const char* v) {
    const char* last = v + std::strlen(v);
    int result;

    auto r = std::from_chars<int>(v, last, result);
    if ( r.ec != std::errc{} || r.ptr != last || result < 0 ) {
        std::fprintf(stderr, "invalid numeric value for %c: %s\n", opt, v);
        std::exit(1);
    }

    return result;
}

/**
 * Default is a small cluster for quick testing.
 */
struct ClusterLayoutOptions {
    bool manager = true;
    int loggers = 1;
    int proxies = 1;
    std::string prefix;
    std::string workers = "1";
    int port = 27760;
    std::string address = "127.0.0.1";
    int metrics_port = 9991;
    std::string metrics_address = "127.0.0.1";
    std::string cluster_node_prefix = "";

    std::string NextPort() {
        auto result = std::to_string(port) + "/tcp";
        ++port;
        return result;
    }

    std::string NextMetricsPort() {
        auto result = std::to_string(metrics_port) + "/tcp";
        // Only increment if not zero: 0/tcp means disabled!
        if ( metrics_port != 0 )
            ++metrics_port;
        return result;
    }

    std::string PrefixedClusterNode(const std::string& name) {
        if ( ! cluster_node_prefix.empty() )
            return cluster_node_prefix + "-" + name;
        return name;
    }
};

class ClusterLayout {
public:
    ClusterLayout() {
        head_lines.emplace_back("# Auto-generated by zeek-cluster-layout-generator");
        head_lines.emplace_back("");

        cluster_nodes_lines.emplace_back("redef Cluster::nodes += {");
    }

    void Done() {
        cluster_nodes_lines.emplace_back("};");
        cluster_nodes_lines.emplace_back("");

        // Multi-host cluster?
        if ( ! cluster_hosts.empty() ) {
            // If this is a multi-host cluster, introduce a Cluster::host variable
            // that contains Cluster::Host records. Do this here for now instead
            // of putting it into the base scripts, because it's a bit of a detail
            // how static clusters work with this generator.
            //
            //  We render something like:
            //
            //    redef Cluster::hosts += {
            //        ["host"] = [$ip=(blocking_lookup_hostname("host") as vector of addr)[0]],
            //    }
            //
            // So that IP addresses of other hosts are resolved during cluster-layout.zeek
            // loading time (ZeekControl instead does that during deploy time in Python).
            cluster_host_lines.emplace_back("module Cluster;");
            cluster_host_lines.emplace_back("");
            cluster_host_lines.emplace_back("type Host: record {");
            cluster_host_lines.emplace_back("\tip: addr;");
            cluster_host_lines.emplace_back("};");
            cluster_host_lines.emplace_back("");
            cluster_host_lines.emplace_back("const hosts: table[string] of Host &redef;");
            cluster_host_lines.emplace_back("");
            cluster_host_lines.emplace_back("redef Cluster::hosts += {");
            for ( const auto& host : cluster_hosts ) {
                std::string line = "\t[\"" + host.name + "\"] = [$ip=";
                if ( host.addr.empty() )
                    line += "(blocking_lookup_hostname(\"" + host.name + "\") as vector of addr)[0]";
                else
                    line += host.addr;
                line += "],";

                cluster_host_lines.emplace_back(std::move(line));
            }
            cluster_host_lines.emplace_back("};");
            cluster_host_lines.emplace_back("");
        }

        if ( total_managers > 0 || total_loggers > 0 )
            tail_lines.emplace_back(total_loggers > 0 ? "redef Cluster::manager_is_logger = F;" :
                                                        "redef Cluster::manager_is_logger = T;");

        tail_lines.emplace_back("");
        tail_lines.emplace_back("@load base/frameworks/telemetry/options");
        tail_lines.emplace_back("");

        tail_lines.emplace_back("# Either use ZEEK_TELEMETRY_LISTEN_ADDRESS env or the value from node$ip.");
        tail_lines.emplace_back("const metrics_address_env = getenv(\"ZEEK_TELEMETRY_LISTEN_ADDRESS\");");
        tail_lines.emplace_back("@if ( |metrics_address_env| > 0 )");
        tail_lines.emplace_back("redef Telemetry::metrics_address = metrics_address_env;");
        tail_lines.emplace_back("@else");
        tail_lines.emplace_back("@if ( Cluster::node in Cluster::nodes )");
        tail_lines.emplace_back("const my_ip = Cluster::nodes[Cluster::node]$ip;");
        tail_lines.emplace_back("const my_ip_str = is_v4_addr(my_ip) ? cat(my_ip) : cat(\"[\", my_ip, \"]\");");
        tail_lines.emplace_back("redef Telemetry::metrics_address = my_ip_str;");
        tail_lines.emplace_back("@endif");
        tail_lines.emplace_back("@endif");
        tail_lines.emplace_back("");
        tail_lines.emplace_back("# Always use the metrics port from the cluster layout for now");
        tail_lines.emplace_back("# This can/may change in the future!");
        tail_lines.emplace_back("redef Telemetry::metrics_port = Cluster::local_node_metrics_port();");
    }


    void AddNode(const std::string& name, const std::string& type, const std::string& ip,
                 const std::string& metrics_port = "", const std::string& port = "", bool has_manager = 1) {
        auto s = "    [\"" + name + "\"] = [$node_type=Cluster::" + type + ", $ip=" + ip;

        if ( ! port.empty() )
            s += ", $p=" + port;

        if ( name != "manager" && has_manager )
            s += ", $manager=\"manager\"";

        if ( ! metrics_port.empty() )
            s += ", $metrics_port=" + metrics_port;

        s += "],";
        cluster_nodes_lines.push_back(std::move(s));

        // Slight hack for the manager_is_logger redef.
        if ( type == "LOGGER" )
            ++total_loggers;

        if ( type == "MANAGER" )
            ++total_managers;
    }

    /**
     * Helper to add workers using the ClusterLayoutOptions.
     *
     * Workers is either just the number of workers, or a "name1:worker1 name2:workers2 name3:workers3"
     * that allows to have worker-{name}-{index} style names.
     */
    void AddWorkers(ClusterLayoutOptions& opts) {
        if ( std::ranges::all_of(opts.workers, [](auto c) { return std::isdigit(c); }) ) {
            int workers = checked_stoi('W', opts.workers.c_str());
            for ( int i = 1; i <= workers; i++ )
                AddNode(opts.PrefixedClusterNode("worker-" + std::to_string(i)), "WORKER", opts.address,
                        opts.NextMetricsPort(), /*port=*/"", /*has_manager=*/opts.manager);
        }
        else {
            // Iterate through the eth1:8,eth2:32,eth3:4 range
            // of interface name followed by number of workers
            // and add workers for each of them.
            std::string workers_optarg = opts.workers;

            // Allow comma or space separated list.
            std::replace(workers_optarg.begin(), workers_optarg.end(), ',', ' ');

            std::cmatch cmatch;
            std::regex re_named_workers("^([-_0-9a-z]+):([0-9]+)( |$)");
            std::string tail;
            const char* p = workers_optarg.c_str();
            while ( std::regex_search(p, cmatch, re_named_workers) ) {
                std::string name = cmatch[1].str();
                std::string nstr = cmatch[2].str();
                tail = cmatch[3].str(); // remember for error checking
                int n = checked_stoi('W', nstr.c_str());

                for ( int i = 1; i <= n; i++ )
                    AddNode(opts.PrefixedClusterNode("worker-" + name + "-" + std::to_string(i)), "WORKER",
                            opts.address, opts.NextMetricsPort(), /*port=*/"", /*has_manager=*/opts.manager);

                p += cmatch.length();
            }

            if ( *p || ! tail.empty() ) {
                std::fprintf(stderr, "-W invalid: '%s'\n", opts.workers.c_str());
                std::exit(1);
            }
        }
    }

    void AddClusterHost(std::string host, std::string addr = "") {
        cluster_hosts.emplace_back(Host{std::move(host), std::move(addr)});
    }

    void Print(std::ostream& out) const {
        for ( const auto& l : head_lines )
            out << l << "\n";
        for ( const auto& l : cluster_host_lines )
            out << l << "\n";
        for ( const auto& l : cluster_nodes_lines )
            out << l << "\n";
        for ( const auto& l : tail_lines )
            out << l << "\n";
    }

    int TotalManagers() const noexcept { return total_managers; }

private:
    int total_loggers = 0;
    int total_managers = 0;
    std::vector<std::string> head_lines;

    struct Host {
        std::string name;
        std::string addr;
    };
    std::vector<Host> cluster_hosts;

    std::vector<std::string> cluster_host_lines;
    std::vector<std::string> cluster_nodes_lines;
    std::vector<std::string> tail_lines;
};

void usage(const char* prog) {
    fprintf(stdout,
            "Usage:\n\n"
            "    %s -L <loggers> -P <proxies> [-W <workers> | -W <iface1>:<workers1>,<iface2>:<workers2>,...] -M 0|1"
            " -a <address> -p <port> -m <metrics_port> -b <metrics_address> -x <cluster_node_prefix> [-o <outfile>]\n"
            "    %s -C /etc/zeek/zeek.conf [-o <outfile>]\n"
            "    %s -C /etc/zeek/cluster [-o <outfile>]\n",
            prog, prog, prog);
}

} // namespace

int main(int argc, char* argv[]) {
    ClusterLayoutOptions opts;

    std::string config_file_or_cluster_dir = "";
    std::string out = "-";
    opterr = 0;

    int c;
    while ( (c = getopt(argc, argv, "hC:L:P:W:M:a:p:m:b:o:x:")) != -1 ) {
        if ( c < 0 )
            break;

        switch ( c ) {
            case 'C': config_file_or_cluster_dir = optarg; break;
            case 'L': opts.loggers = checked_stoi(c, optarg); break;
            case 'P': opts.proxies = checked_stoi(c, optarg); break;
            case 'W': opts.workers = optarg; break;
            case 'M': {
                int manager = checked_stoi(c, optarg);
                if ( manager != 0 && manager != 1 ) {
                    std::fprintf(stderr, "flag -M expects 0 or 1\n");
                    std::exit(1);
                }
                opts.manager = manager == 1 ? true : false;
                break;
            }
            case 'a': opts.address = optarg; break;
            case 'p': opts.port = checked_stoi(c, optarg); break;
            case 'm': opts.metrics_port = checked_stoi(c, optarg); break;
            case 'b': opts.metrics_address = optarg; break;
            case 'o': out = optarg; break;
            case 'x': opts.cluster_node_prefix = optarg; break;
            case 'h': usage(argv[0]); std::exit(0);
            case '?': std::fprintf(stderr, "unrecognized flag %c\n", optopt); std::exit(1);
        }
    }

    ClusterLayout layout;

    if ( config_file_or_cluster_dir.empty() ) {
        if ( opts.manager )
            layout.AddNode("manager", "MANAGER", opts.address, opts.NextMetricsPort(), opts.NextPort());

        for ( int i = 1; i <= opts.loggers; i++ ) {
            layout.AddNode(opts.PrefixedClusterNode("logger-" + std::to_string(i)), "LOGGER", opts.address,
                           opts.NextMetricsPort(), opts.NextPort(), /*has_manager=*/opts.manager);
        }

        for ( int i = 1; i <= opts.proxies; i++ )
            layout.AddNode(opts.PrefixedClusterNode("proxy-" + std::to_string(i)), "PROXY", opts.address,
                           opts.NextMetricsPort(), opts.NextPort(), /*has_manager=*/opts.manager);

        layout.AddWorkers(opts);
    }
    else {
        // -C was used to select either a configuration file, or a cluster directory.
        std::vector<std::filesystem::path> fnames;

        if ( std::filesystem::is_directory(config_file_or_cluster_dir) ) {
            // If the -C argument is a directory, parse all configuration files from it and
            // assemble a cluster layout. Configs need to strictly match <hostname>.zeek.conf
            // within the directory.
            //
            // We could validate a few things here, like IP address an name conflicts,
            // but for now this should do, and technically overlapping IPs are okay
            // if the ports differ, so there's that. But IP/port overlap among processes
            // would be problematic.
            std::regex re_conf("[a-z0-9][-a-z0-9_]*\\.zeek\\.conf$");

            auto dir = std::filesystem::path(config_file_or_cluster_dir);
            if ( dir.filename().empty() ) // strip trailing slash, can only be one.
                dir = dir.parent_path();


            for ( const auto& fname : std::filesystem::directory_iterator(dir) ) {
                if ( std::regex_match(fname.path().filename().string(), re_conf) )
                    fnames.emplace_back(fname);
            }

            if ( fnames.empty() ) {
                std::fprintf(stderr, "no <hostname>.zeek.conf files in %s found\n", config_file_or_cluster_dir.c_str());
                std::exit(1);
            }
        }
        else if ( std::filesystem::is_regular_file(config_file_or_cluster_dir) ) {
            // If it's just a file, do the same.
            fnames.emplace_back(config_file_or_cluster_dir);
        }
        else {
            std::fprintf(stderr, "neither file nor directory: %s\n", config_file_or_cluster_dir.c_str());
            std::exit(1);
        }


        // Make processing order predictable.
        std::sort(fnames.begin(), fnames.end());

        for ( const auto& fname : fnames ) {
            const auto absfname = std::filesystem::absolute(fname);
            auto config = zeek::detail::parse_config("/dev/null", absfname);

            // Failed to load the config file? Strange. Just exit.
            if ( ! config.Exists() || ! config.IsValid() ) {
                std::fprintf(stderr, "invalid config %s\n", absfname.string().c_str());
                for ( const auto& err : config.Errors() )
                    std::fprintf(stderr, "%s\n", err.c_str());

                std::exit(1);
            }

            int port = config.ClusterPort();
            auto next_port = [&port]() { return std::to_string(port++) + "/tcp"; };

            int metrics_port = config.MetricsPort();
            auto next_metrics_port = [&metrics_port]() {
                auto port_str = std::to_string(metrics_port) + "/tcp";
                if ( metrics_port != 0 )
                    metrics_port++;
                return port_str;
            };

            // If there is no cluster address available, introduce
            // the Cluster::hosts table to look it up dynamically.
            std::string addr;
            if ( config.HasFilenameHost() ) {
                layout.AddClusterHost(config.FilenameHost(), config.ClusterAddress());
                addr = "hosts[\"" + config.FilenameHost() + "\"]$ip";
            }
            else {
                // This should default to 127.0.0.1 if -C /path/to/zeek.conf was used.
                addr = config.ClusterAddress();
            }

            // The manager in a cluster is always named manager, no prefix here.
            if ( config.Manager() )
                layout.AddNode("manager", "MANAGER", addr, next_metrics_port(), next_port());

            for ( int i = 1; i <= config.Loggers(); i++ ) {
                layout.AddNode(config.PrefixedClusterNode("logger-" + std::to_string(i)), "LOGGER", addr,
                               next_metrics_port(), next_port());
            }

            for ( int i = 1; i <= config.Proxies(); i++ )
                layout.AddNode(config.PrefixedClusterNode("proxy-" + std::to_string(i)), "PROXY", addr,
                               next_metrics_port(), next_port());

            // Go through all interface configurations and add the required workers.
            for ( const auto& iwc : config.InterfaceWorkerConfigs() ) {
                for ( int i = 1; i <= iwc.Workers(); i++ ) {
                    layout.AddNode(config.PrefixedClusterNode(iwc.FullWorkerName(i)), "WORKER", addr,
                                   next_metrics_port());
                }
            }
        }

        if ( layout.TotalManagers() > 1 ) {
            fprintf(stderr, "too many manager nodes (%d) from %s\n", layout.TotalManagers(),
                    config_file_or_cluster_dir.c_str());
            std::exit(1);
        }
    }

    layout.Done();

    // stdout print?
    if ( out == "-" ) {
        layout.Print(std::cout);
        return 0;
    }

    std::ofstream os(out, std::ios::trunc);
    if ( ! os.is_open() ) {
        std::error_code ec{errno, std::system_category()};
        std::fprintf(stderr, "failed to open %s: %s\n", out.c_str(), ec.message().c_str());
        std::exit(1);
    }

    layout.Print(os);
    return 0;
}
