/// @file test/profile/party_profile.cpp
/// @brief Workload for the (2+1) party protocols.
///
/// Spawns p0/p1/p2 the same way party_bench does. Each role reports CPU time
/// and the bytes and frames of the flow itself. The repeat barrier is not
/// included in those counters. `wall_ms` includes process startup.
///
///   profile_party --list
///   profile_party --suite core --repeat 5 --warmup 1
///   profile_party --suite extreme
///   profile_party --suite gadget --repeat 3 --warmup 1
///   profile_party --tag "beaver,bench" --repeat 3
///   profile_party --case beaver_dot_n32
///
/// Profile-guided build. The party binaries must be rebuilt with the same
/// mode, because the protocol code lives in those processes:
///   cmake -S test -B build-pgo -DLIBDPF_PGO=generate -DCMAKE_BUILD_TYPE=Release
///   cmake --build build-pgo --target profile_party p0 p1 p2
///   build-pgo/bin/profile_party --suite all --repeat 4 --warmup 1
///   cmake -S test -B build-pgo -DLIBDPF_PGO=use -DCMAKE_BUILD_TYPE=Release
///   cmake --build build-pgo --target profile_party p0 p1 p2

#include "cases.hpp"
#include "registry.hpp"
#include "spawn.hpp"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <string>
#include <map>
#include <sstream>
#include <vector>

namespace
{

const char * const kCore[] = {
    "beaver_dot_n32",
    "beaver_product_100_200",
    "beaver_stream_n128",
    "beaver_horner_d4",
    "beaver_batch_one_round",
    "beaver_xor_mux",
    "dcf_gt",
    "dpf_point_2a_7",
    "geneval_point",
    "geneval_arith_point",
    "blocked_dcf_point",
    "ds_key_agrees",
    "grotto_prefix_horner",
    "verifiable_honest_point",
};

const char * const kExtreme[] = {
    "beaver_dot_n128",
    "beaver_stream_n2048",
    "dcf_dense_gt",
    "dcf_blocked_interval_ip",
    "grotto_signed_prefix_dense",
    "grotto_offset_horner_d3",
    "grotto_offset_horner_multipiece",
    "grotto_geneval_offset_horner",
    "geneval_interval_dense",
    "geneval_cmp_dense_gt",
    "ds_cmp_many_points",
    "recent_offset_poly",
    "recent_offset_jet",
    "recent_offset_twist",
    "recent_offset_repr",
    "recent_dist_dpf3_point",
    "wildcard_single_leaf",
};

const char * const kGadget[] = {
    "arith_proj_m5",
    "arith_proj_m17",
    "arith_proj_m64",
    "arith_mul_p5",
    "arith_mul_p7",
    "arith_mul_p11",
    "arith_thresh_b8",
    "arith_thresh_b16",
    "arith_chain_mul4",
    "yao_if_4_2",
    "yao_if_16_8",
    "yao_onehot_k4",
    "yao_onehot_k8",
    "flute_d2",
    "flute_d4",
    "flute_d8",
    "flute_d4_o8",
    "shuffle_n16",
    "shuffle_n64",
    "shuffle_n256",
};

const char kUsage[] =
    "profile_party [--list] [--slice S ...] [--tier std|heavy|all|smoke]\n"
    "              [--suite core|extreme|gadget|all] [--tag TAGS] [--case NAME ...]\n"
    "              [--repeat N=3] [--warmup W=1]\n"
    "Default run is the core suite. --suite gadget is the word-garbling,\n"
    "stacked-Yao, FLUTE, and hidden-shuffle battery on the party mesh.\n"
    "--slice/--tier select the cost-matrix catalog instead.\n"
    "--list with no filter prints every bench flow.\n"
    "heavy: wildcard_single_leaf, beaver_stream_n512/n2048, dcf_full_*, *_domain.\n"
    "smoke: flows that are not benchable (rejects and negative tests).\n";

bool starts_with(const std::string & s, const char * prefix)
{
    const auto n = std::char_traits<char>::length(prefix);
    return s.size() >= n && s.compare(0, n, prefix) == 0;
}

std::string slice_of(const std::string & name)
{
    if (starts_with(name, "arith_"))
        return "arith-garble";
    if (starts_with(name, "yao_"))
        return "yao-stack";
    if (starts_with(name, "flute_"))
        return "flute";
    if (starts_with(name, "shuffle_"))
        return "shuffle";
    if (starts_with(name, "beaver_dot"))
        return "beaver-dot";
    if (starts_with(name, "beaver_stream"))
        return "beaver-stream";
    if (starts_with(name, "beaver_scale"))
        return "beaver-scale";
    if (starts_with(name, "beaver_horner") || name == "beaver_sign_horner")
        return "beaver-horner";
    if (starts_with(name, "beaver_product") || starts_with(name, "beaver_mul")
        || starts_with(name, "beaver_chained"))
        return "beaver-product";
    if (starts_with(name, "beaver_poly") || name == "beaver_like_terms"
        || starts_with(name, "beaver_nested") || starts_with(name, "beaver_tuple")
        || starts_with(name, "beaver_mixed") || starts_with(name, "beaver_factored"))
        return "beaver-poly";
    if (starts_with(name, "beaver_"))
        return "beaver-misc";
    if (starts_with(name, "dcf_full"))
        return "dcf-full";
    if (starts_with(name, "dcf_dense"))
        return "dcf-dense";
    if (starts_with(name, "dcf_blocked") || starts_with(name, "blocked_"))
        return "blocked";
    if (starts_with(name, "dcf_"))
        return "dcf";
    if (starts_with(name, "geneval_cmp"))
        return "geneval-cmp";
    if (starts_with(name, "geneval_"))
        return "geneval";
    if (starts_with(name, "grotto_"))
        return "grotto-net";
    if (starts_with(name, "ds_"))
        return "ds";
    if (starts_with(name, "recent_offset") || starts_with(name, "recent_ring")
        || starts_with(name, "recent_closed") || starts_with(name, "recent_exact"))
        return "recent-grotto";
    if (starts_with(name, "recent_dpf3") || starts_with(name, "recent_dist_dpf3"))
        return "dpf3";
    if (starts_with(name, "recent_"))
        return "recent";
    if (starts_with(name, "wildcard_"))
        return "wildcard";
    if (starts_with(name, "verifiable_"))
        return "verifiable";
    if (starts_with(name, "dpf_"))
        return "dpf-point";
    if (starts_with(name, "carry_") || starts_with(name, "extractable_"))
        return "auth";
    if (starts_with(name, "cov_paint") || starts_with(name, "cov_idcf")
        || name == "cov_eq_at" || name == "cov_ccmp")
        return "coverage-cmp";
    if (starts_with(name, "cov_geneval") || name == "cov_idpf_at"
        || name == "cov_eval_inner_product")
        return "coverage-eval";
    if (starts_with(name, "cov_"))
        return "coverage";
    return "other";
}

std::string tier_of(const std::string & name, bool)
{
    const bool domain = name.size() >= 7
        && name.compare(name.size() - 7, 7, "_domain") == 0;
    const bool dpf3_walk = name.find("dpf3") != std::string::npos
        && (name.find("update") != std::string::npos
            || name.find("proof") != std::string::npos
            || name.find("blocked") != std::string::npos);
    if (name == "wildcard_single_leaf"
        || name == "beaver_stream_n512"
        || name == "beaver_stream_n2048"
        || starts_with(name, "dcf_full_")
        || domain
        || dpf3_walk
        || starts_with(name, "recent_dist_dpf3"))
        return "heavy";
    const bool negative = name.find("fail") != std::string::npos
        || name.find("tamper") != std::string::npos
        || name.find("reject") != std::string::npos
        || name.find("wrong") != std::string::npos
        || name.find("bad_use") != std::string::npos;
    if (negative)
        return "smoke";
    return "std";
}

std::map<std::string, std::string> fields_of(const std::string & line)
{
    std::map<std::string, std::string> out;
    std::istringstream in(line);
    std::string tok;
    while (in >> tok)
    {
        const auto eq = tok.find('=');
        if (eq == std::string::npos)
            continue;
        out.emplace(tok.substr(0, eq), tok.substr(eq + 1));
    }
    return out;
}

std::string field(const std::map<std::string, std::string> & m, const char * key)
{
    const auto it = m.find(key);
    return it == m.end() ? "0" : it->second;
}

}  // namespace

int main(int argc, char ** argv)
{
    dpf::party::register_all_flows();

    std::string suite = "core";
    std::string tag;
    std::string tier;
    std::vector<std::string> slices;
    std::vector<std::string> case_names;
    std::uint64_t repeat = 3;
    std::uint64_t warmup = 1;
    bool list_only = false;
    bool suite_set = false;
    bool tier_set = false;

    for (int i = 1; i < argc; ++i)
    {
        const std::string a = argv[i];
        auto need = [&](const char * flag) {
            if (i + 1 >= argc)
            {
                std::cerr << "missing value for " << flag << "\n";
                std::exit(2);
            }
            return std::string(argv[++i]);
        };
        if (a == "--list")
            list_only = true;
        else if (a == "--suite")
        {
            suite = need("--suite");
            suite_set = true;
        }
        else if (a == "--tag")
            tag = need("--tag");
        else if (a == "--case")
            case_names.push_back(need("--case"));
        else if (a == "--slice")
            slices.push_back(need("--slice"));
        else if (a == "--tier")
        {
            tier = need("--tier");
            tier_set = true;
        }
        else if (a == "--repeat")
            repeat = std::stoull(need("--repeat"));
        else if (a == "--warmup")
            warmup = std::stoull(need("--warmup"));
        else if (a == "--help" || a == "-h")
        {
            std::cout << kUsage;
            return 0;
        }
        else
        {
            std::cerr << "unknown argument: " << a << "\n" << kUsage;
            return 2;
        }
    }
    if (repeat == 0)
        repeat = 1;

    const bool catalog_mode = tier_set || !slices.empty()
        || (list_only && !suite_set && tag.empty() && case_names.empty());

    struct picked
    {
        std::string name;
        std::string slice;
        std::string tier;
    };
    std::vector<picked> picked_flows;

    auto take_catalog = [&](const std::string & want_tier) {
        for (const auto * f : dpf::party::select_flows(""))
        {
            const std::string slice = slice_of(f->name);
            const std::string flow_tier = tier_of(f->name, f->bench);
            if (want_tier == "std" && flow_tier != "std")
                continue;
            if (want_tier == "heavy" && flow_tier != "heavy")
                continue;
            if (want_tier == "smoke" && flow_tier != "smoke")
                continue;
            if (want_tier == "all" && flow_tier == "smoke")
                continue;
            if (!slices.empty()
                && std::find(slices.begin(), slices.end(), slice) == slices.end())
                continue;
            if (!case_names.empty()
                && std::find(case_names.begin(), case_names.end(), f->name) == case_names.end())
                continue;
            picked_flows.push_back(picked{f->name, slice, flow_tier});
        }
    };

    if (catalog_mode)
    {
        const std::string want = tier_set ? tier : (list_only ? "all" : "std");
        if (want != "std" && want != "heavy" && want != "all" && want != "smoke")
        {
            std::cerr << "unknown tier: " << want << "\n" << kUsage;
            return 2;
        }
        take_catalog(want);
        for (const auto & name : case_names)
        {
            bool found = false;
            for (const auto & row : picked_flows)
                found = found || row.name == name;
            if (!found)
            {
                std::cerr << "unknown flow: " << name << "\n";
                return 2;
            }
        }
        if (picked_flows.empty())
        {
            std::cerr << "no flows match the requested slice/tier/case\n";
            return 2;
        }
    }
    else
    {
        std::vector<std::string> names;
        if (!case_names.empty())
            names = std::move(case_names);
        else if (!tag.empty())
        {
            for (const auto * f : dpf::party::select_flows(tag))
                names.emplace_back(f->name);
            if (names.empty())
            {
                std::cerr << "no flows match tag: " << tag << "\n";
                return 2;
            }
        }
        else if (suite == "core" || suite == "all")
        {
            for (const char * n : kCore)
                names.emplace_back(n);
            if (suite == "all")
            {
                for (const char * n : kExtreme)
                    names.emplace_back(n);
                for (const char * n : kGadget)
                    names.emplace_back(n);
            }
        }
        else if (suite == "extreme")
        {
            for (const char * n : kExtreme)
                names.emplace_back(n);
        }
        else if (suite == "gadget")
        {
            for (const char * n : kGadget)
                names.emplace_back(n);
        }
        else
        {
            std::cerr << "unknown suite: " << suite << "\n" << kUsage;
            return 2;
        }

        for (const auto & name : names)
        {
            const auto * f = dpf::party::find_flow(name);
            if (!f)
            {
                std::cerr << "unknown flow: " << name << "\n";
                return 2;
            }
            picked_flows.push_back(picked{name, slice_of(name), tier_of(name, f->bench)});
        }
    }

    if (list_only)
    {
        std::cout << "family\tslice\tcase\titems\ttier\n";
        for (const auto & row : picked_flows)
            std::cout << "party\t" << row.slice << '\t' << row.name
                      << "\t1\t" << row.tier << '\n';
        return 0;
    }

    dpf::party::spawn_opts opts;
    opts.repeat = repeat;
    opts.warmup = warmup;
    opts.metrics = true;

    std::cout
        << "family\tslice\ttier\tflow\trole\twall_ms\tavg_ns\tmin_ns\tmax_ns\t"
        << "bytes_sent\tbytes_recv\tframes_sent\tframes_recv\t"
        << "bytes_recv_from_p2\tbytes_recv_from_peer\t"
        << "bytes_sent_to_p2\tbytes_sent_to_peer\t"
        << "payload_sent\tpayload_recv\twire_overhead\trounds\tprg_evals\t"
        << "avg_bytes_sent\tavg_frames_sent\trc\n";

    int fails = 0;
    for (const auto & row : picked_flows)
    {
        const auto result = dpf::party::spawn_trio_flow(row.name, opts);
        if (result.metrics_lines.empty())
        {
            std::cout << "party\t" << row.slice << '\t' << row.tier << '\t'
                      << row.name << "\t-\t" << result.wall_ms
                      << "\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t0\t"
                      << result.rc[0] << "," << result.rc[1] << "," << result.rc[2]
                      << "\n";
        }
        for (const auto & line : result.metrics_lines)
        {
            const auto f = fields_of(line);
            const auto role = field(f, "role");
            const int rc = role == "p0" ? result.rc[0]
                : role == "p1" ? result.rc[1]
                : role == "p2" ? result.rc[2] : 0;
            std::cout << "party\t" << row.slice << '\t' << row.tier << '\t'
                      << row.name << '\t' << role << '\t' << result.wall_ms << '\t'
                      << field(f, "avg_ns") << '\t'
                      << field(f, "min_ns") << '\t'
                      << field(f, "max_ns") << '\t'
                      << field(f, "bytes_sent") << '\t'
                      << field(f, "bytes_recv") << '\t'
                      << field(f, "frames_sent") << '\t'
                      << field(f, "frames_recv") << '\t'
                      << field(f, "bytes_recv_from_p2") << '\t'
                      << field(f, "bytes_recv_from_peer") << '\t'
                      << field(f, "bytes_sent_to_p2") << '\t'
                      << field(f, "bytes_sent_to_peer") << '\t'
                      << field(f, "payload_sent") << '\t'
                      << field(f, "payload_recv") << '\t'
                      << field(f, "wire_overhead") << '\t'
                      << field(f, "rounds") << '\t'
                      << field(f, "prg_evals") << '\t'
                      << field(f, "avg_bytes_sent") << '\t'
                      << field(f, "avg_frames_sent") << '\t'
                      << rc << '\n';
        }
        if (result.rc[0] || result.rc[1] || result.rc[2])
            ++fails;
    }
    return fails ? 1 : 0;
}