#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "dpf/app_flow.hpp"
#include "dpf/app_plans.hpp"
#include "dpf/bench_cells.hpp"
#include "dpf/experiment.hpp"
#include "dpf/party_runner.hpp"
#include "dpf/run_log.hpp"

// Every cell is a compose plan driven by `run_parties` on one `run_config`:
// the party mesh over in-process async memory, unix sockets, TCP mux, parallel
// TCP, or SCTP. `memory` and `stream` are the older paired sinks; this harness
// uses the mesh, so those two names run as async. Every `run_config` key is a
// flag (`--transport=mux --lanes=4 --window=262144 --warmup=2 --trials=9`) or
// the matching `DPF_*` variable; flags win. Lanes default to 1 here. Each cell
// runs `warmup` untimed and `trials` timed repetitions and records every
// party's trial times, party 0's and the slowest party's medians, the
// configuration, and party 0's wire counters in the CSVs. Every repetition
// draws from streams derived from the cell's master, so replaying the master
// replays the cell.
//
//   c++ -std=c++17 -march=native -pthread -I include -I thirdparty \
//       examples/applications/experiment_bench.cpp -lsctp
//   DPF_EXPERIMENT_DIR=/tmp/libdpf_bench ./a.out --transport=mux --trials=5
//
// The first block is the DPF plans. The second block is the work that is not
// a key: word gadgets, stacked branches, short tables, and a hidden reorder
// of a column the parties already share.

namespace
{

dpf::app::run_config mesh_config(int argc, char ** argv, std::string & replaced)
{
    dpf::app::run_config cfg;
    cfg.n_lanes = 1;
    cfg.merge_env();
    const auto rest = cfg.apply_args(argc, argv);
    if (!rest.empty())
        throw std::invalid_argument("unexpected argument '" + rest.front()
            + "' (flags are --key=value)");
    if (cfg.kind == dpf::net::transport::memory_sink
        || cfg.kind == dpf::net::transport::memory_stream)
    {
        std::cout << "transport "
                  << dpf::net::transport_name(cfg.kind)
                  << " is a paired sink; the battery uses the party mesh (async)\n";
        replaced = dpf::net::transport_name(cfg.kind);
        cfg.kind = dpf::net::transport::async_memory;
    }
    return cfg;
}

int measure_plans(const char * name, std::vector<dpf::protocol::plan> plans,
    std::uint64_t run_id, const std::string & dir, const dpf::app::run_config & cfg,
    dpf::protocol::cell_fn cell)
{
    if (plans.empty() || plans[0].rounds() == 0)
    {
        std::cerr << name << " has no exchange rounds\n";
        return 1;
    }
    dpf::experiment ex(name, "p0");
    ex.set_run_id(run_id);
    ex.ingest_plan(plans[0]);
    ex.set_config(cfg.describe());
    dpf::app::parties_result result;
    const std::size_t total = cfg.warmup + std::max<std::size_t>(1, cfg.trials);
    try
    {
        for (std::size_t t = 0; t < total; ++t)
        {
            std::vector<dpf::app::party_values> values(plans.size());
            const bool last = t + 1 == total;
            result = dpf::app::run_parties(plans, values, {}, cfg,
                last ? &ex : nullptr, cell, &ex);
            if (t >= cfg.warmup)
                ex.add_trial(result.party0_wall_ns, result.party_wall_ns);
        }
    }
    catch (const std::exception & err)
    {
        std::cerr << name << " flow: " << err.what() << "\n";
        return 1;
    }
    const auto wire = result.wire.empty() ? dpf::net::stream_stats{} : result.wire[0];
    dpf::experiment::wire_counts w;
    w.bytes_out = wire.bytes_out;
    w.bytes_in = wire.bytes_in;
    w.payload_out = wire.payload_out;
    w.payload_in = wire.payload_in;
    w.frames_out = wire.frames_out;
    w.frames_in = wire.frames_in;
    w.write_calls = wire.write_calls;
    ex.set_wire(w);
    ex.write_csv(dir);
    std::cout << name << " parties=" << plans.size()
              << " run_id=" << run_id
              << " rounds=" << ex.interactive_rounds()
              << " bytes=" << ex.plan_bytes_out()
              << " wire_out=" << wire.bytes_out
              << " wire_in=" << wire.bytes_in
              << " payload_out=" << wire.payload_out
              << " median_ns=" << ex.median_trial_ns()
              << " slowest_median_ns=" << ex.slowest_median_ns()
              << " trials=" << ex.trials().size()
              << " prg_evals=" << ex.prg_evals()
              << " seed=" << ex.seed_hex() << "\n";
    return 0;
}

}  // namespace

int main(int argc, char ** argv)
{
    const char * env = std::getenv("DPF_EXPERIMENT_DIR");
    const std::string dir = (env && env[0] != '\0') ? env
                                                    : "/tmp/libdpf_experiment_bench";
    dpf::app::run_config cfg;
    std::string replaced;
    try
    {
        cfg = mesh_config(argc, argv, replaced);
        dpf::app::start_logging(cfg);
    }
    catch (const std::exception & err)
    {
        std::cerr << "experiment_bench: " << err.what() << "\n";
        return 2;
    }
    if (!replaced.empty())
        DPF_LOG(warning, "config.override").kv("key", "transport")
            .kv("requested", replaced).kv("used", "async")
            .kv("detail", "paired sinks cannot carry the party mesh");
    std::cout << cfg.summary() << "\n";

    struct named
    {
        const char * name;
        int parties;
        dpf::protocol::plan (*make)(std::size_t party);
    };
    const named dpf_plans[] = {
        {"keyword_pir", 2, [](std::size_t p) {
             return dpf::protocol::keyword_pir_compose_plan(p, 8);
         }},
        {"express", 2, [](std::size_t p) {
             return dpf::protocol::mailbox_write_fused_plan(p);
         }},
        {"subleq", 2, [](std::size_t p) {
             return dpf::protocol::subleq_instruction_plan(p);
         }},
        {"pika", 2, [](std::size_t p) {
             return dpf::protocol::pika_lookup_plan(p);
         }},
        {"duoram3", 2, [](std::size_t p) {
             return dpf::protocol::duoram_update_plan(p);
         }},
        {"poplar", 2, [](std::size_t p) {
             return dpf::protocol::poplar_prefix_plan(p);
         }},
        {"poplar_fan4", 2, [](std::size_t p) {
             return dpf::protocol::poplar_prefix_fan_plan(p, 4);
         }},
        {"ledger23", 2, [](std::size_t p) {
             return dpf::protocol::ledger23_append_plan(p);
         }},
        {"bitmore", 2, [](std::size_t p) {
             return dpf::protocol::bitmore_fan_plan(p);
         }},
        {"floram", 2, [](std::size_t p) {
             return dpf::protocol::floram_ds_plan(p);
         }},
        {"fss_point", 2, [](std::size_t p) {
             return dpf::protocol::fss_point_plan(p);
         }},
        {"fss_cmp", 2, [](std::size_t p) {
             return dpf::protocol::fss_cmp_plan(p);
         }},
        {"range_count", 2, [](std::size_t p) {
             return dpf::protocol::range_count_plan(p);
         }},
        {"psi_cuckoo", 2, [](std::size_t p) {
             return dpf::protocol::psi_cuckoo_plan(p, {0, 2, 5, 7});
         }},
        {"idpf_agg", 2, [](std::size_t p) {
             return dpf::protocol::idpf_agg_plan(p, 8);
         }},
    };

    std::uint64_t run_id = 0;
    for (const auto & row : dpf_plans)
    {
        std::vector<dpf::protocol::plan> plans;
        plans.reserve(static_cast<std::size_t>(row.parties));
        for (int p = 0; p < row.parties; ++p)
            plans.push_back(row.make(static_cast<std::size_t>(p)));
        if (int rc = measure_plans(row.name, std::move(plans), run_id++, dir, cfg,
                nullptr))
            return rc;
    }

    for (const auto & cell : dpf::bench::battery())
    {
        std::vector<dpf::protocol::plan> plans;
        plans.reserve(static_cast<std::size_t>(cell.parties));
        for (int p = 0; p < cell.parties; ++p)
            plans.push_back(dpf::bench::plan_for(static_cast<std::size_t>(p), cell.id));
        if (int rc = measure_plans(cell.name, std::move(plans), run_id++, dir, cfg,
                &dpf::bench::run_cell))
            return rc;
    }
    std::cout << "wrote CSVs under " << dir << "\n";
    return 0;
}