#include "dpf/random.hpp"
#include <tuple>

#include <gtest/gtest.h>

#include <array>
#include <cstdint>
#include <cstring>
#include <limits>
#include <set>
#include <thread>
#include <vector>

#include <sys/wait.h>
#include <unistd.h>

#include "dpf/bit.hpp"
#include "dpf/xor_wrapper.hpp"
#include "simde/simde/x86/avx2.h"

namespace dpf_rng_test
{
std::uint64_t sample_from_other_tu();
}

namespace
{

struct HookState
{
    int calls = 0;
    std::size_t last_n = 0;
    unsigned char byte = 0xA5;
    bool use_typed = false;
    std::int32_t typed = 1;
    std::vector<std::size_t> sizes;
    const unsigned char * script = nullptr;
    std::size_t script_i = 0;
};

HookState g_hook;

void reset_hook()
{
    g_hook = HookState{};
    dpf::detail::uniform_bytes_hook = nullptr;
}

void test_hook(void * dst, std::size_t n)
{
    ++g_hook.calls;
    g_hook.last_n = n;
    g_hook.sizes.push_back(n);
    if (g_hook.script != nullptr)
    {
        std::memcpy(dst, g_hook.script + g_hook.script_i, n);
        g_hook.script_i += n;
        return;
    }
    if (g_hook.use_typed && n == sizeof(g_hook.typed))
    {
        std::memcpy(dst, &g_hook.typed, n);
        return;
    }
    std::memset(dst, g_hook.byte, n);
}

void fill_byte(void * dst, std::size_t n, unsigned char byte)
{
    std::memset(dst, byte, n);
}

template <typename T>
T load_bytes(unsigned char byte)
{
    T value;
    fill_byte(&value, sizeof(value), byte);
    return value;
}

bool same_bytes(const void * a, const void * b, std::size_t n)
{
    return std::memcmp(a, b, n) == 0;
}

template <typename T>
void expect_reconstruct_unsigned(T secret)
{
    auto [share0, share1] = dpf::additively_share(secret);
    EXPECT_EQ(dpf::reconstruct(share0, share1), secret);
}

template <typename T>
void expect_reconstruct_signed(T secret)
{
    auto [share0, share1] = dpf::additively_share(secret);
    using U = std::make_unsigned_t<T>;
    EXPECT_EQ(static_cast<U>(dpf::reconstruct(share0, share1)),
        static_cast<U>(secret));
}

}  // namespace

class RandomTest : public ::testing::Test
{
  protected:
    void TearDown() override
    {
        reset_hook();
    }
};

TEST_F(RandomTest, UniformFillReturnsTheSameObject)
{
    std::uint32_t value = 0;
    EXPECT_EQ(&dpf::uniform_fill(value), &value);
}

TEST_F(RandomTest, HookOverwritesEveryByte)
{
    dpf::detail::uniform_bytes_hook = test_hook;
    g_hook.byte = 0xA5;

    struct Padded
    {
        std::uint8_t a;
        std::uint32_t b;
    };
    Padded padded = dpf::uniform_sample<Padded>();
    EXPECT_EQ(g_hook.calls, 1);
    EXPECT_EQ(g_hook.last_n, sizeof(Padded));
    auto expected = load_bytes<Padded>(0xA5);
    EXPECT_TRUE(same_bytes(&padded, &expected, sizeof(Padded)));
}

TEST_F(RandomTest, HookSeesTheObjectSize)
{
    dpf::detail::uniform_bytes_hook = test_hook;
    g_hook.byte = 0x3C;

    auto byte = dpf::uniform_sample<std::uint8_t>();
    auto block = dpf::uniform_sample<simde__m128i>();
    std::array<unsigned char, 32> wide{};
    dpf::uniform_fill(wide);

    EXPECT_EQ(byte, static_cast<std::uint8_t>(0x3C));
    EXPECT_EQ(g_hook.sizes, (std::vector<std::size_t>{
        1u, sizeof(simde__m128i), 32u}));
    auto expected_block = load_bytes<simde__m128i>(0x3C);
    EXPECT_TRUE(same_bytes(&block, &expected_block, sizeof(block)));
    for (unsigned char b : wide)
        EXPECT_EQ(b, static_cast<unsigned char>(0x3C));
}

TEST_F(RandomTest, ClearingTheHookReadsTheDeviceAgain)
{
    dpf::detail::uniform_bytes_hook = [](void * dst, std::size_t n)
    {
        fill_byte(dst, n, 0x11);
    };
    EXPECT_EQ(dpf::uniform_sample<std::uint32_t>(), 0x11111111u);
    dpf::detail::uniform_bytes_hook = nullptr;

    std::set<std::array<unsigned char, 16>> draws;
    for (int i = 0; i < 8; ++i)
    {
        std::array<unsigned char, 16> sample{};
        dpf::uniform_fill(sample);
        draws.insert(sample);
        bool nonzero = false;
        for (unsigned char b : sample)
            nonzero = nonzero || b != 0;
        EXPECT_TRUE(nonzero);
    }
    EXPECT_EQ(draws.size(), 8u);
}

TEST_F(RandomTest, HookIsThreadLocal)
{
    dpf::detail::uniform_bytes_hook = [](void * dst, std::size_t n)
    {
        fill_byte(dst, n, 0x11);
    };

    std::uint32_t other = 0;
    std::thread worker([&]()
    {
        dpf::detail::uniform_bytes_hook = [](void * dst, std::size_t n)
        {
            fill_byte(dst, n, 0x22);
        };
        other = dpf::uniform_sample<std::uint32_t>();
        dpf::detail::uniform_bytes_hook = nullptr;
    });
    worker.join();

    EXPECT_EQ(dpf::uniform_sample<std::uint32_t>(), 0x11111111u);
    EXPECT_EQ(other, 0x22222222u);
}

TEST_F(RandomTest, BooleanSamplesUseOnlyTheLowBit)
{
    const unsigned char raw[] = {0x00, 0x02, 0xFF};
    g_hook.script = raw;
    dpf::detail::uniform_bytes_hook = test_hook;

    EXPECT_FALSE(dpf::uniform_sample<bool>());
    EXPECT_FALSE(dpf::uniform_sample<bool>());
    EXPECT_TRUE(dpf::uniform_sample<bool>());

    g_hook.script_i = 0;
    EXPECT_EQ(dpf::uniform_sample<dpf::bit>(), dpf::bit::zero);
    EXPECT_EQ(dpf::uniform_sample<dpf::bit>(), dpf::bit::zero);
    EXPECT_EQ(dpf::uniform_sample<dpf::bit>(), dpf::bit::one);
    EXPECT_EQ(g_hook.calls, 6);
}

TEST_F(RandomTest, DeviceBooleanSamplesAreBothValues)
{
    bool saw_false = false;
    bool saw_true = false;
    bool saw_zero = false;
    bool saw_one = false;
    for (int i = 0; i < 64; ++i)
    {
        bool bit = dpf::uniform_sample<bool>();
        saw_false = saw_false || !bit;
        saw_true = saw_true || bit;
        dpf::bit as_bit = dpf::uniform_sample<dpf::bit>();
        EXPECT_TRUE(as_bit == dpf::bit::zero || as_bit == dpf::bit::one);
        saw_zero = saw_zero || as_bit == dpf::bit::zero;
        saw_one = saw_one || as_bit == dpf::bit::one;
    }
    EXPECT_TRUE(saw_false);
    EXPECT_TRUE(saw_true);
    EXPECT_TRUE(saw_zero);
    EXPECT_TRUE(saw_one);
}

TEST_F(RandomTest, AdditiveSharesReconstructUnsigned)
{
    expect_reconstruct_unsigned<std::uint8_t>(0);
    expect_reconstruct_unsigned<std::uint8_t>(1);
    expect_reconstruct_unsigned<std::uint8_t>(255);
    expect_reconstruct_unsigned<std::uint32_t>(0);
    expect_reconstruct_unsigned<std::uint32_t>(1);
    expect_reconstruct_unsigned<std::uint32_t>(0xFFFFFFFFu);
    expect_reconstruct_unsigned<std::uint64_t>(0);
    expect_reconstruct_unsigned<std::uint64_t>(
        std::numeric_limits<std::uint64_t>::max());

    dpf::detail::uniform_bytes_hook = test_hook;
    g_hook.byte = 0xFF;
    auto [share0, share1] = dpf::additively_share<std::uint32_t>(0);
    EXPECT_EQ(g_hook.calls, 1);
    EXPECT_EQ(share0.raw(), 0xFFFFFFFFu);
    EXPECT_EQ(share1.raw(), 1u);
    EXPECT_EQ(dpf::reconstruct(share0, share1), 0u);
}

TEST_F(RandomTest, AdditiveSharesReconstructSignedExtremes)
{
    dpf::detail::uniform_bytes_hook = test_hook;
    g_hook.use_typed = true;
    g_hook.typed = 1;

    constexpr auto secret = std::numeric_limits<std::int32_t>::min();
    auto [share0, share1] = dpf::additively_share(secret);
    EXPECT_EQ(share0.raw(), 1);
    using U = std::uint32_t;
    EXPECT_EQ(static_cast<U>(dpf::reconstruct(share0, share1)),
        static_cast<U>(secret));

    dpf::detail::uniform_bytes_hook = nullptr;
    expect_reconstruct_signed<std::int8_t>(std::numeric_limits<std::int8_t>::min());
    expect_reconstruct_signed<std::int8_t>(std::numeric_limits<std::int8_t>::max());
    expect_reconstruct_signed<std::int8_t>(-1);
    expect_reconstruct_signed<std::int32_t>(secret);
    expect_reconstruct_signed<std::int32_t>(std::numeric_limits<std::int32_t>::max());
    expect_reconstruct_signed<std::int32_t>(-1);
    expect_reconstruct_signed<std::int32_t>(0);
    expect_reconstruct_signed<std::int64_t>(std::numeric_limits<std::int64_t>::min());
    expect_reconstruct_signed<std::int64_t>(std::numeric_limits<std::int64_t>::max());
}

TEST_F(RandomTest, AdditiveSharesReconstructBitAndXorGroup)
{
    for (dpf::bit secret : {dpf::bit::zero, dpf::bit::one})
    {
        auto [share0, share1] = dpf::additively_share(secret);
        EXPECT_TRUE(share0.raw() == dpf::bit::zero || share0.raw() == dpf::bit::one);
        EXPECT_TRUE(share1.raw() == dpf::bit::zero || share1.raw() == dpf::bit::one);
        EXPECT_EQ(dpf::reconstruct(share0, share1), secret);
    }

    dpf::xor_wrapper<std::uint32_t> secret{0xA5A5A5A5u};
    auto [share0, share1] = dpf::additively_share(secret);
    EXPECT_EQ(dpf::reconstruct(share0, share1), secret);
}

TEST_F(RandomTest, AdditiveShareDrawsOnce)
{
    dpf::detail::uniform_bytes_hook = test_hook;
    g_hook.byte = 0x01;
    auto shares = dpf::additively_share<std::uint32_t>(0x10u);
    EXPECT_EQ(g_hook.calls, 1);
    EXPECT_EQ(shares.first.raw(), 0x01010101u);
    EXPECT_EQ(dpf::reconstruct(shares.first, shares.second), 0x10u);
}

TEST_F(RandomTest, ForkedProcessesDoNotRepeatEntropy)
{
    // A buffered stdio read of the device copies the unread buffer into the
    // child. Prime the generator, then compare the next draw on each side.
    dpf::uniform_sample<std::uint64_t>();

    int fds[2];
    ASSERT_EQ(::pipe(fds), 0);
    pid_t pid = ::fork();
    ASSERT_NE(pid, -1);
    if (pid == 0)
    {
        ::close(fds[0]);
        auto sample = dpf::uniform_sample<std::array<unsigned char, 32>>();
        ssize_t wrote = ::write(fds[1], sample.data(), sample.size());
        ::_exit(wrote == static_cast<ssize_t>(sample.size()) ? 0 : 1);
    }

    ::close(fds[1]);
    auto parent = dpf::uniform_sample<std::array<unsigned char, 32>>();
    std::array<unsigned char, 32> child{};
    std::size_t got = 0;
    while (got < child.size())
    {
        ssize_t n = ::read(fds[0], child.data() + got, child.size() - got);
        if (n <= 0)
            break;
        got += static_cast<std::size_t>(n);
    }
    ::close(fds[0]);
    int status = 0;
    ASSERT_EQ(::waitpid(pid, &status, 0), pid);
    ASSERT_TRUE(WIFEXITED(status));
    ASSERT_EQ(WEXITSTATUS(status), 0);
    ASSERT_EQ(got, child.size());
    EXPECT_FALSE(same_bytes(parent.data(), child.data(), parent.size()));
}

TEST_F(RandomTest, ConcurrentReadsStayWellFormed)
{
    constexpr int threads = 4;
    constexpr int draws = 128;
    std::vector<std::vector<std::uint64_t>> results(threads);
    std::vector<std::thread> workers;
    workers.reserve(threads);
    for (int t = 0; t < threads; ++t)
    {
        workers.emplace_back([&, t]()
        {
            results[t].reserve(draws);
            for (int i = 0; i < draws; ++i)
                results[t].push_back(dpf::uniform_sample<std::uint64_t>());
        });
    }
    for (auto & worker : workers)
        worker.join();

    std::set<std::uint64_t> unique;
    for (const auto & row : results)
    {
        EXPECT_EQ(row.size(), static_cast<std::size_t>(draws));
        unique.insert(row.begin(), row.end());
    }
    EXPECT_EQ(unique.size(), static_cast<std::size_t>(threads * draws));
}

TEST_F(RandomTest, OtherTranslationUnitCanSample)
{
    auto here = dpf::uniform_sample<std::uint64_t>();
    auto there = dpf_rng_test::sample_from_other_tu();
    EXPECT_NE(here, there);
}

TEST_F(RandomTest, DeviceLowBitTakesBothValues)
{
    // Doerner–Shelat pad bits are `uniform_sample<unsigned char>() & 1`.
    bool saw0 = false;
    bool saw1 = false;
    for (int i = 0; i < 64; ++i)
    {
        auto bit = static_cast<unsigned>(dpf::uniform_sample<unsigned char>() & 1u);
        saw0 = saw0 || bit == 0u;
        saw1 = saw1 || bit == 1u;
    }
    EXPECT_TRUE(saw0);
    EXPECT_TRUE(saw1);
}