#include <gtest/gtest.h>

#include <array>
#include <cstdint>
#include <cstring>
#include <exception>
#include <filesystem>
#include <string>
#include <thread>
#include <unistd.h>
#include <utility>
#include <vector>

#include "aes_mmo_ref.hpp"
#include "dpf/net/trio.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"
#include "dpf/yao.hpp"
#include "dpf/yao_aes.hpp"
#include "dpf/verifiable.hpp"
#include "dpf/yao_share.hpp"

namespace
{

using dpf::yao::bit;
using dpf::yao::bits8;
using dpf::yao::netlist;

void bytes_to_bits(const std::uint8_t * bytes, int nbytes, std::uint8_t * bits)
{
    for (int i = 0; i < nbytes; ++i)
        for (int k = 0; k < 8; ++k)
            bits[i * 8 + k] = static_cast<std::uint8_t>((bytes[i] >> (7 - k)) & 1u);
}

void bits_to_bytes(const std::uint8_t * bits, int nbytes, std::uint8_t * bytes)
{
    for (int i = 0; i < nbytes; ++i)
    {
        bytes[i] = 0;
        for (int k = 0; k < 8; ++k)
            bytes[i] = static_cast<std::uint8_t>(bytes[i] | (bits[i * 8 + k] << (7 - k)));
    }
}

std::uint8_t pack_byte(const std::uint8_t * bits)
{
    std::uint8_t b = 0;
    bits_to_bytes(bits, 1, &b);
    return b;
}

netlist sbox_netlist()
{
    netlist n;
    const bits8 in = dpf::yao::shared_byte(n);
    dpf::yao::out_byte(n, dpf::yao::sbox(n, in));
    return n;
}

void expand_key(const std::uint8_t key[16], std::uint8_t rk[11][16])
{
    std::uint8_t w[176];
    std::memcpy(w, key, 16);
    int n = 16;
    std::uint8_t rcon = 1;
    while (n < 176)
    {
        std::uint8_t t[4];
        std::memcpy(t, w + n - 4, 4);
        if (n % 16 == 0)
        {
            const std::uint8_t tmp = t[0];
            t[0] = t[1];
            t[1] = t[2];
            t[2] = t[3];
            t[3] = tmp;
            for (int i = 0; i < 4; ++i)
                t[i] = dpf::party::aes_ref::sbox_at(t[i]);
            t[0] = static_cast<std::uint8_t>(t[0] ^ rcon);
            rcon = dpf::party::aes_ref::xtime(rcon);
        }
        for (int i = 0; i < 4; ++i)
        {
            w[n] = static_cast<std::uint8_t>(w[n - 16] ^ t[i]);
            ++n;
        }
    }
    std::memcpy(rk, w, 176);
}

void encrypt_aes128(std::uint8_t s[16], const std::uint8_t rk[11][16])
{
    using namespace dpf::party::aes_ref;
    add_round_key(s, rk[0]);
    for (int r = 1; r < 10; ++r)
    {
        sub_bytes(s);
        shift_rows(s);
        mix_columns(s);
        add_round_key(s, rk[r]);
    }
    sub_bytes(s);
    shift_rows(s);
    add_round_key(s, rk[10]);
}

std::vector<std::uint8_t> reconstruct(std::vector<std::uint8_t> a,
    const std::vector<std::uint8_t> & b)
{
    EXPECT_EQ(a.size(), b.size());
    for (std::size_t i = 0; i < a.size(); ++i)
        a[i] = static_cast<std::uint8_t>(a[i] ^ b[i]);
    return a;
}

template <typename Fn0, typename Fn1>
void run_pair(Fn0 && fn0, Fn1 && fn1)
{
    const auto dir = std::filesystem::temp_directory_path()
        / ("libdpf_yao_" + std::to_string(::getpid()));
    std::filesystem::create_directories(dir);
    std::exception_ptr ep0;
    std::exception_ptr ep1;
    std::thread t0([&] {
        try
        {
            auto net = dpf::net::trio::connect_pair(dpf::net::role::p0, dir.string());
            fn0(net);
        }
        catch (...)
        {
            ep0 = std::current_exception();
        }
    });
    std::thread t1([&] {
        try
        {
            auto net = dpf::net::trio::connect_pair(dpf::net::role::p1, dir.string());
            fn1(net);
        }
        catch (...)
        {
            ep1 = std::current_exception();
        }
    });
    t0.join();
    t1.join();
    std::filesystem::remove_all(dir);
    if (ep0)
        std::rethrow_exception(ep0);
    if (ep1)
        std::rethrow_exception(ep1);
}

}  // namespace

TEST(Yao, AndXorNot)
{
    netlist n;
    const bit a = n.shared_in();
    const bit b = n.shared_in();
    const bit x = n.xor_(a, b);
    const bit y = n.and_(a, b);
    const bit z = n.not_(a);
    const bit w = n.xnor_(a, b);
    n.out(x);
    n.out(y);
    n.out(z);
    n.out(w);
    n.out(n.xor_public(y, 1));
    EXPECT_EQ(n.n_and(), 1u);
    for (int av = 0; av < 2; ++av)
    {
        for (int bv = 0; bv < 2; ++bv)
        {
            const std::uint8_t in[2] = {
                static_cast<std::uint8_t>(av), static_cast<std::uint8_t>(bv)};
            const auto plain = dpf::yao::eval_plain(n, in);
            const auto got = dpf::yao::eval_local(n, in);
            const std::uint8_t expect[5] = {
                static_cast<std::uint8_t>(av ^ bv),
                static_cast<std::uint8_t>(av & bv),
                static_cast<std::uint8_t>(av ^ 1),
                static_cast<std::uint8_t>(1 ^ av ^ bv),
                static_cast<std::uint8_t>((av & bv) ^ 1)};
            EXPECT_EQ(plain, (std::vector<std::uint8_t>(expect, expect + 5)));
            EXPECT_EQ(got, plain);
        }
    }
}

TEST(Yao, SboxPlainAndGarbled)
{
    const netlist n = sbox_netlist();
    EXPECT_EQ(n.n_and(), static_cast<std::uint32_t>(aes_bp::and_count));
    for (int x = 0; x < 256; ++x)
    {
        const auto byte = static_cast<std::uint8_t>(x);
        std::uint8_t bits[8];
        bytes_to_bits(&byte, 1, bits);
        const auto plain = dpf::yao::eval_plain(n, bits);
        const auto garbled = dpf::yao::eval_local(n, bits);
        EXPECT_EQ(pack_byte(plain.data()), dpf::party::aes_ref::sbox_at(byte)) << x;
        EXPECT_EQ(garbled, plain) << x;
    }
}

TEST(Yao, SboxShares)
{
    const netlist n = sbox_netlist();
    const std::uint8_t byte = 0x53;
    std::uint8_t semantic[8];
    bytes_to_bits(&byte, 1, semantic);
    std::uint8_t p0[8], p1[8];
    for (int i = 0; i < 8; ++i)
    {
        p0[i] = static_cast<std::uint8_t>(dpf::uniform_sample<unsigned char>() & 1u);
        p1[i] = static_cast<std::uint8_t>(semantic[i] ^ p0[i]);
    }
    auto [a, b] = dpf::yao::eval_pair(n, p0, p1);
    EXPECT_EQ(pack_byte(reconstruct(a, b).data()), dpf::party::aes_ref::sbox_at(byte));
}

TEST(Yao, Aes128Nist)
{
    const std::uint8_t key[16] = {
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
    const std::uint8_t pt[16] = {
        0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff};
    const std::uint8_t ct[16] = {
        0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30,
        0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a};

    std::uint8_t rk[11][16];
    expand_key(key, rk);
    std::uint8_t state[16];
    std::memcpy(state, pt, 16);
    encrypt_aes128(state, rk);
    EXPECT_EQ(std::vector<std::uint8_t>(state, state + 16),
        std::vector<std::uint8_t>(ct, ct + 16));

    const netlist n = dpf::yao::aes128_netlist();
    EXPECT_EQ(n.n_and(), 6400u);
    std::uint8_t bits[256];
    bytes_to_bits(pt, 16, bits);
    bytes_to_bits(key, 16, bits + 128);
    const auto plain = dpf::yao::eval_plain(n, bits);
    std::uint8_t got[16];
    bits_to_bytes(plain.data(), 16, got);
    EXPECT_EQ(std::vector<std::uint8_t>(got, got + 16),
        std::vector<std::uint8_t>(ct, ct + 16));

    const auto garbled = dpf::yao::eval_local(n, bits);
    bits_to_bytes(garbled.data(), 16, got);
    EXPECT_EQ(std::vector<std::uint8_t>(got, got + 16),
        std::vector<std::uint8_t>(ct, ct + 16));
}

TEST(Yao, AesMmoMatchesPrg)
{
    const std::uint32_t positions[] = {0u, 1u, 0x01020304u};
    const std::uint8_t msg[16] = {
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
        0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c};
    simde__m128i block{};
    std::memcpy(&block, msg, 16);
    std::uint8_t bits[128];
    dpf::yao::block_to_bits(bits, block);
    for (std::uint32_t pos : positions)
    {
        const auto expect = dpf::prg::aes128::eval(block, pos);
        const netlist n = dpf::yao::aes_mmo_netlist(pos);
        EXPECT_EQ(n.n_and(), 5120u);
        const auto got_bits = dpf::yao::eval_local(n, bits);
        const auto got = dpf::yao::bits_to_block(got_bits.data());
        std::uint8_t a[16], b[16];
        std::memcpy(a, &expect, 16);
        std::memcpy(b, &got, 16);
        EXPECT_EQ(std::vector<std::uint8_t>(a, a + 16),
            std::vector<std::uint8_t>(b, b + 16)) << pos;
    }
}

TEST(Yao, SessionSboxAndReuse)
{
    const netlist gates = [] {
        netlist n;
        const bit a = n.priv_in(0);
        const bit b = n.priv_in(1);
        const bit c = n.shared_in();
        n.out(n.xor_(n.and_(a, b), c));
        n.out(n.not_(c));
        return n;
    }();
    const netlist priv_xor = [] {
        netlist n;
        const bit a = n.priv_in(0);
        const bit b = n.priv_in(0);
        n.out(n.xor_(a, b));
        return n;
    }();
    const netlist box = sbox_netlist();
    const std::uint8_t byte = 0xa7;
    std::uint8_t semantic[8];
    bytes_to_bits(&byte, 1, semantic);
    std::uint8_t sp0[8], sp1[8];
    for (int i = 0; i < 8; ++i)
    {
        sp0[i] = static_cast<std::uint8_t>(i & 1u);
        sp1[i] = static_cast<std::uint8_t>(semantic[i] ^ sp0[i]);
    }

    std::vector<std::uint8_t> g0, g1, s0, s1, x0, x1;
    run_pair(
        [&](dpf::net::trio & net) {
            dpf::yao::session s;
            const std::uint8_t xin[2] = {1, 1};
            x0 = s.eval(0, priv_xor, xin, net.to(dpf::net::role::p1));
            const std::uint8_t in0[3] = {1, 0, 1};
            g0 = s.eval(0, gates, in0, net.to(dpf::net::role::p1));
            const std::uint8_t in1[3] = {1, 0, 0};
            auto second = s.eval(0, gates, in1, net.to(dpf::net::role::p1));
            g0.insert(g0.end(), second.begin(), second.end());
            s0 = s.eval(0, box, sp0, net.to(dpf::net::role::p1));
        },
        [&](dpf::net::trio & net) {
            dpf::yao::session s;
            const std::uint8_t xin[2] = {0, 0};
            x1 = s.eval(1, priv_xor, xin, net.to(dpf::net::role::p0));
            const std::uint8_t in0[3] = {0, 1, 0};
            g1 = s.eval(1, gates, in0, net.to(dpf::net::role::p0));
            const std::uint8_t in1[3] = {0, 0, 1};
            auto second = s.eval(1, gates, in1, net.to(dpf::net::role::p0));
            g1.insert(g1.end(), second.begin(), second.end());
            s1 = s.eval(1, box, sp1, net.to(dpf::net::role::p0));
        });

    EXPECT_EQ(reconstruct(x0, x1), std::vector<std::uint8_t>{0});
    const auto opened = reconstruct(g0, g1);
    const std::uint8_t expect_gates[4] = {0, 0, 1, 0};
    EXPECT_EQ(std::vector<std::uint8_t>(opened.begin(), opened.begin() + 4),
        std::vector<std::uint8_t>(expect_gates, expect_gates + 4));
    EXPECT_EQ(pack_byte(reconstruct(s0, s1).data()),
        dpf::party::aes_ref::sbox_at(byte));
}

TEST(Yao, LeafBitsOnSession)
{
    const std::uint32_t beta = 0x6bu;
    const auto s0 = dpf::subtractive_share<std::uint32_t, 0>::from_raw(beta);
    const auto s1 = dpf::subtractive_share<std::uint32_t, 1>::from_raw(0);
    const auto [p0, p1] = dpf::yao::b2y(s0, s1, 8);
    const dpf::yao::netlist n = [] {
        dpf::yao::netlist c;
        dpf::yao::bit in[8];
        for (int i = 0; i < 8; ++i)
            in[i] = c.shared_in();
        c.out(c.and_(in[0], in[1]));
        return c;
    }();
    std::vector<std::uint8_t> o0, o1;
    bool locked = false;
    run_pair(
        [&](dpf::net::trio & net) {
            dpf::yao::session s;
            auto & link = net.to(dpf::net::role::p1);
            o0 = s.eval(0, n, p0.data(), link);
            try
            {
                s.eval(1, n, p0.data(), link);
            }
            catch (const std::logic_error &)
            {
                locked = true;
            }
        },
        [&](dpf::net::trio & net) {
            dpf::yao::session s;
            o1 = s.eval(1, n, p1.data(), net.to(dpf::net::role::p0));
        });
    EXPECT_TRUE(locked);
    ASSERT_EQ(o0.size(), 1u);
    EXPECT_EQ(static_cast<std::uint8_t>(o0[0] ^ o1[0]), 1u);
    const auto [z0, z1] = dpf::yao::y2b<std::uint32_t>(o0, o1, 1);
    EXPECT_EQ(dpf::reconstruct(z0, z1), 1u);
}

TEST(Yao, CorrectionLevelMatchesMakeCs)
{
    const simde__m128i s0 = dpf::uniform_sample<simde__m128i>();
    const simde__m128i s1 = dpf::uniform_sample<simde__m128i>();
    const std::uint64_t prefix = 0x0123456789abcdefull;
    const std::uint64_t share0 = 0x1111222233334444ull;
    const std::uint64_t share1 = prefix ^ share0;
    const std::size_t level = 0x10ffu;
    std::uint8_t in0[dpf::yao::correction_level_in_bits];
    std::uint8_t in1[dpf::yao::correction_level_in_bits];
    dpf::yao::pack_correction_level(0, s0, share0, level, in0);
    dpf::yao::pack_correction_level(1, s1, share1, level, in1);
    const auto nl = dpf::yao::correction_level_netlist();
    EXPECT_EQ(nl.n_and(), dpf::yao::correction_level_ands);
    EXPECT_EQ(nl.n_in(), static_cast<std::uint32_t>(dpf::yao::correction_level_in_bits));
    auto [a, b] = dpf::yao::eval_pair(nl, in0, in1);
    ASSERT_EQ(a.size(), dpf::yao::correction_level_out_bits);
    const auto expect = dpf::detail::vdpf::make_cs(level, prefix, s0, s1);
    for (int lane = 0; lane < 4; ++lane)
    {
        std::uint8_t bits[128];
        for (int i = 0; i < 128; ++i)
            bits[i] = static_cast<std::uint8_t>(
                a[static_cast<std::size_t>(lane * 128 + i)]
                ^ b[static_cast<std::size_t>(lane * 128 + i)]);
        const auto got = dpf::yao::bits_to_block(bits);
        EXPECT_EQ(std::memcmp(&got, &expect[static_cast<std::size_t>(lane)], 16), 0)
            << lane;
    }
}