#include <gtest/gtest.h>

#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <string>
#include <thread>
#include <unistd.h>
#include <vector>

#include "dpf.hpp"
#include "dpf/app_flow.hpp"
#include "dpf/app_plans.hpp"
#include "dpf/beaver.hpp"
#include "dpf/buffered_prg.hpp"
#include "dpf/doerner_shelat.hpp"
#include "dpf/experiment.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"

namespace
{

std::string tmp_dir(const char * tag)
{
    const std::string dir =
        std::string("/tmp/libdpf_ex_") + tag + "_" + std::to_string(::getpid());
    (void)::system(("rm -rf " + dir + " && mkdir -p " + dir).c_str());
    return dir;
}

std::size_t count_lines(const std::string & path)
{
    std::ifstream in(path);
    std::size_t n = 0;
    std::string line;
    while (std::getline(in, line))
        ++n;
    return n;
}

TEST(Experiment, ReplayMatchesMakeDpfRoots)
{
    dpf::experiment::master_seed master{};
    std::uint8_t alpha = 7;
    using node_t = dpf::prg::aes128::block_type;
    node_t root0a{}, root0b{}, root1a{}, root1b{};
    {
        dpf::experiment ex("dpf");
        master = ex.seed();
        auto [k0, k1] = dpf::make_dpf(alpha, std::uint64_t{1});
        root0a = k0.root();
        root1a = k1.root();
    }
    {
        auto ex = dpf::experiment::replay("dpf", master);
        auto [k0, k1] = dpf::make_dpf(alpha, std::uint64_t{1});
        root0b = k0.root();
        root1b = k1.root();
    }
    EXPECT_EQ(std::memcmp(&root0a, &root0b, sizeof(root0a)), 0);
    EXPECT_EQ(std::memcmp(&root1a, &root1b, sizeof(root1a)), 0);
}

TEST(Experiment, NestedRaiiRestoresOuterHook)
{
    dpf::experiment outer("outer");
    const auto outer_seed = outer.seed();
    const auto a = dpf::uniform_sample<std::uint64_t>();
    {
        dpf::experiment inner("inner");
        EXPECT_NE(inner.seed(), outer_seed);
        (void)dpf::uniform_sample<std::uint64_t>();
    }
    // Outer stream continues where it left off (inner restored the hook).
    const auto b = dpf::uniform_sample<std::uint64_t>();
    auto again = dpf::experiment::replay("outer", outer_seed);
    EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), a);
    EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), b);
}

TEST(Experiment, NoteSeedFromConstructors)
{
    dpf::experiment ex("ctors");
    EXPECT_TRUE(ex.has_seed_named("master"));
    {
        dpf::beavers::oracle<std::uint64_t> o(16);
        EXPECT_TRUE(ex.has_seed_named("beavers::oracle"));
        EXPECT_TRUE(ex.has_seed_named("lane_table"));
        (void)o;
    }
    {
        dpf::randomness::buffered_prg<dpf::prg::aes128, std::uint64_t> prg(8);
        EXPECT_TRUE(ex.has_seed_named("buffered_prg"));
        (void)prg;
    }
    {
        dpf::prg_pad_rng<> pad;
        EXPECT_TRUE(ex.has_seed_named("prg_pad_rng"));
        (void)pad;
    }
    {
        dpf::pseudorandom_root_sampler<dpf::prg::aes128> roots;
        EXPECT_TRUE(ex.has_seed_named("pseudorandom_root_sampler"));
        (void)roots;
    }
    ex.note_seed("custom", std::uint32_t{0xdeadbeefu});
    EXPECT_TRUE(ex.has_seed_named("custom"));
    EXPECT_GE(ex.seed_count(), 6u);
}

TEST(Experiment, RandomBytesCounterTracksUniformFill)
{
    dpf::reset_random_bytes_count();
    EXPECT_EQ(dpf::random_bytes_count(), 0u);
    dpf::experiment ex("bytes");
    // ctor resets after drawing the master.
    EXPECT_EQ(dpf::random_bytes_count(), 0u);
    (void)dpf::uniform_sample<std::uint64_t>();
    EXPECT_EQ(dpf::random_bytes_count(), 8u);
    (void)dpf::uniform_sample<std::uint32_t>();
    EXPECT_EQ(dpf::random_bytes_count(), 12u);
}

TEST(Experiment, IngestPlanCriticalPathAndEdges)
{
    dpf::experiment ex("plan");
    auto p = dpf::protocol::fss_point_plan(0, 8, 16);
    ex.ingest_plan(p);
    EXPECT_EQ(ex.interactive_rounds(), 8u);
    EXPECT_EQ(ex.dag_depth(), p.waves());
    EXPECT_EQ(ex.plan_bytes_out(), 128u);
    EXPECT_EQ(ex.plan_edge_bytes(dpf::protocol::edge_channel::peer), 128u);
    EXPECT_EQ(ex.plan_edge_bytes(dpf::protocol::edge_channel::dealer), 0u);
    EXPECT_GT(ex.critical_path_length(), 0u);
}

TEST(Experiment, MeasurePlanProbeAndTiming)
{
    auto plan = dpf::protocol::mailbox_write_fused_plan(0);
    auto ex = dpf::app::measure_plan("mailbox", plan);
    EXPECT_EQ(ex.interactive_rounds(), 8u);
    EXPECT_EQ(ex.rounds().size(), 8u);
    std::size_t sum_out = 0, sum_in = 0;
    for (const auto & r : ex.rounds())
    {
        EXPECT_EQ(r.channel, dpf::protocol::edge_channel::peer);
        sum_out += r.bytes_out;
        sum_in += r.bytes_in;
    }
    EXPECT_EQ(sum_out, ex.bytes_out());
    EXPECT_EQ(sum_in, ex.bytes_in());
    EXPECT_EQ(ex.edge_bytes_out(dpf::protocol::edge_channel::peer),
        ex.bytes_out());
    EXPECT_GT(ex.wall_ns(), 0u);
    EXPECT_GT(ex.prg_evals(), 0u);  // walk kernels expand
    EXPECT_TRUE(ex.has_seed_named("master"));
}

TEST(Experiment, MeasureClientServersTwoRounds)
{
    auto plan = dpf::protocol::n_server_pir_plan(0, 2, 64, 8);
    const auto slots = plan.slot_bytes_all();
    ASSERT_EQ(slots.size(), 2u);
    auto ex = dpf::app::measure_plan("pir2", plan);
    EXPECT_EQ(ex.interactive_rounds(), 2u);
    ASSERT_EQ(ex.rounds().size(), 2u);
    EXPECT_EQ(ex.rounds()[0].bytes_out, slots[0]);
    EXPECT_EQ(ex.rounds()[1].bytes_out, slots[1]);
    EXPECT_EQ(ex.plan_bytes_out(), slots[0] + slots[1]);
}

TEST(Experiment, WriteCsvAllFilesAndAppend)
{
    const auto dir = tmp_dir("csv");
    auto plan = dpf::protocol::fss_cmp_plan(0);
    auto ex = dpf::app::measure_plan("cmp", plan);
    ex.set_run_id(1);
    ex.write_csv(dir);
    EXPECT_EQ(count_lines(dir + "/summary.csv"), 2u);  // header + row
    EXPECT_GE(count_lines(dir + "/rounds.csv"), 2u);
    EXPECT_GE(count_lines(dir + "/edges.csv"), 2u);
    EXPECT_GE(count_lines(dir + "/seeds.csv"), 2u);
    EXPECT_GE(count_lines(dir + "/critical_path.csv"), 2u);

    ex.set_run_id(2);
    ex.write_csv(dir);
    EXPECT_EQ(count_lines(dir + "/summary.csv"), 3u);  // append

    std::ifstream seeds(dir + "/seeds.csv");
    std::string line;
    bool saw_master = false;
    while (std::getline(seeds, line))
        if (line.find("master") != std::string::npos)
            saw_master = true;
    EXPECT_TRUE(saw_master);
    (void)::system(("rm -rf " + dir).c_str());
}

TEST(Experiment, RunMeasuredEnvCsv)
{
    const auto dir = tmp_dir("env");
    ASSERT_EQ(::setenv("DPF_EXPERIMENT_DIR", dir.c_str(), 1), 0);
    const int rc = dpf::app::run_measured("env_csv",
        dpf::protocol::keyword_pir_compose_plan(0, 8), 2);
    EXPECT_EQ(rc, 0);
    EXPECT_GE(count_lines(dir + "/summary.csv"), 2u);
    ASSERT_EQ(::unsetenv("DPF_EXPERIMENT_DIR"), 0);
    (void)::system(("rm -rf " + dir).c_str());
}

TEST(Experiment, ProbeAbsentDriveStillWorks)
{
    // exercise_plan without experiment must not require clocks/probes.
    auto plan = dpf::protocol::range_count_plan(0);
    const auto cost = dpf::app::exercise_plan(plan);
    EXPECT_EQ(cost.rounds, 8u);
    EXPECT_EQ(cost.bytes, 128u);
}

TEST(Experiment, UninstallRestoresSystemEntropy)
{
    dpf::experiment::master_seed master{};
    {
        dpf::experiment ex("tmp");
        master = ex.seed();
        (void)dpf::uniform_sample<std::uint64_t>();
    }
    // After destroy, draws are fresh system entropy (not the replay stream).
    // Just ensure we can sample without a hook and without throwing.
    const auto x = dpf::uniform_sample<std::uint64_t>();
    const auto y = dpf::uniform_sample<std::uint64_t>();
    (void)x;
    (void)y;
    auto replayed = dpf::experiment::replay("tmp", master);
    // Replay still works after a window of system entropy.
    (void)dpf::uniform_sample<std::uint64_t>();
}

TEST(Experiment, ThreadsDoNotCrossStreams)
{
    constexpr int n = 32;
    std::vector<std::uint64_t> a(n), b(n);
    dpf::experiment::master_seed sa{}, sb{};
    std::thread t0([&] {
        dpf::experiment ex("A", "p0");
        sa = ex.seed();
        for (int i = 0; i < n; ++i)
            a[i] = dpf::uniform_sample<std::uint64_t>();
    });
    std::thread t1([&] {
        dpf::experiment ex("B", "p1");
        sb = ex.seed();
        for (int i = 0; i < n; ++i)
            b[i] = dpf::uniform_sample<std::uint64_t>();
    });
    t0.join();
    t1.join();
    EXPECT_NE(sa, sb);
    auto ra = dpf::experiment::replay("A", sa);
    for (int i = 0; i < n; ++i)
        EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), a[i]);
    auto rb = dpf::experiment::replay("B", sb);
    for (int i = 0; i < n; ++i)
        EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), b[i]);
}

TEST(Experiment, MoveTransfersActiveContext)
{
    dpf::experiment::master_seed master{};
    std::uint64_t first = 0;
    {
        dpf::experiment ex("move");
        master = ex.seed();
        first = dpf::uniform_sample<std::uint64_t>();
        dpf::experiment moved = std::move(ex);
        EXPECT_EQ(moved.seed(), master);
        const auto second = dpf::uniform_sample<std::uint64_t>();
        auto again = dpf::experiment::replay("move", master);
        EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), first);
        EXPECT_EQ(dpf::uniform_sample<std::uint64_t>(), second);
    }
}

TEST(Experiment, BeginEndTimingWithoutProbe)
{
    dpf::experiment ex("timing");
    ex.begin_timing();
    volatile std::uint64_t sink = 0;
    for (int i = 0; i < 1000; ++i)
        sink += dpf::uniform_sample<std::uint64_t>();
    (void)sink;
    ex.end_timing();
    EXPECT_GT(ex.wall_ns(), 0u);
    EXPECT_GE(ex.random_bytes(), 8000u);
}

}  // namespace