#include <gtest/gtest.h>

#include <cstdint>
#include <string>
#include <type_traits>

#include "dpf.hpp"
#include "dpf/json.hpp"

static_assert(NLOHMANN_JSON_VERSION_MAJOR == 3
    && NLOHMANN_JSON_VERSION_MINOR == 12
    && NLOHMANN_JSON_VERSION_PATCH == 0,
    "JSON tests build against nlohmann 3.12.0");

namespace
{

// Round-trip a comparison-only multi-level key through JSON and confirm the
// (public) comparison channel still reconstructs after deserialization.
TEST(IncrementalJsonTest, CmpKeyRoundTrips)
{
    const uint32_t alpha = 0x00abcdefu;
    const uint64_t yt = 42u;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lt(yt));
    using KT0 = std::decay_t<decltype(k0)>;
    using KT1 = std::decay_t<decltype(k1)>;
    static_assert(KT0::is_multilevel, "cmp key must be multi-level");
    ASSERT_EQ(KT0::num_outputs, 0u);

    const std::string s0 = dpf::json::to_json(k0);
    const std::string s1 = dpf::json::to_json(k1);

    auto r0 = dpf::json::from_json<KT0>(s0);
    auto r1 = dpf::json::from_json<KT1>(s1);

    const uint64_t mask = k0.cmp().mask;
    EXPECT_EQ(r0.cmp().nbits, k0.cmp().nbits);
    EXPECT_EQ(r0.cmp().mask, k0.cmp().mask);
    EXPECT_EQ(r0.cw_last(), k0.cw_last());
    EXPECT_EQ(r0.cmp_addend(), k0.cmp_addend());

    auto recon_cmp = [&](uint32_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, r0, q), dpf::eval_point(dpf::cmp, r1, q)) & mask;
    };
    EXPECT_EQ(recon_cmp(alpha - 1u), yt);
    EXPECT_EQ(recon_cmp(0u), yt);
    EXPECT_EQ(recon_cmp(alpha), 0u);
    EXPECT_EQ(recon_cmp(alpha + 1u), 0u);
}

// A geq comparison (inverted path-sum) must also survive the round-trip.
TEST(IncrementalJsonTest, CmpGeqRoundTrips)
{
    const uint32_t alpha = 100u;
    const uint64_t yt = 5u, yf = 9u;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::geq(yt, yf));
    using KT0 = std::decay_t<decltype(k0)>;
    using KT1 = std::decay_t<decltype(k1)>;

    auto r0 = dpf::json::from_json<KT0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<KT1>(dpf::json::to_json(k1));

    const uint64_t mask = k0.cmp().mask;
    auto r = [&](uint32_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, r0, q), dpf::eval_point(dpf::cmp, r1, q)) & mask;
    };
    EXPECT_EQ(r(99u), yf);
    EXPECT_EQ(r(100u), yt);
    EXPECT_EQ(r(101u), yt);
}

// A single-output classic key, including its leaf share.
TEST(IncrementalJsonTest, ClassicPointRoundTrips)
{
    const uint32_t alpha = 0x00abcdefu;
    const uint32_t y = 99u;
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;

    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));

    EXPECT_EQ(static_cast<uint32_t>(dpf::reconstruct(
        *dpf::eval_point(r0, alpha), *dpf::eval_point(r1, alpha))), y);
    EXPECT_EQ(static_cast<uint32_t>(dpf::reconstruct(
        *dpf::eval_point(r0, alpha ^ 1u), *dpf::eval_point(r1, alpha ^ 1u))), 0u);
}

// Multi-lane outputs are stored as leaf blocks, not truncated integers.
TEST(IncrementalJsonTest, VecOutputRoundTrips)
{
    using out_t = dpf::vec<std::uint32_t, 4>;
    const std::uint16_t alpha = 0x1234;
    out_t y;
    y[0] = 1;
    y[1] = 0xffffffffu;
    y[2] = 7;
    y[3] = 100;
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));
    auto at = [&](std::uint16_t x) {
        return dpf::reconstruct(*dpf::eval_point(r0, x), *dpf::eval_point(r1, x));
    };
    EXPECT_EQ(at(alpha), y);
    EXPECT_EQ(at(0), out_t{});
    EXPECT_EQ(at(static_cast<std::uint16_t>(alpha + 1)), out_t{});
}

// `at<>` keys carry leaf outputs the comparison-only path used to drop.
TEST(IncrementalJsonTest, AtOutputRoundTrips)
{
    const uint32_t x = 0x00abcdefu;
    auto [k0, k1] = dpf::make_dpf(x, dpf::at<10>(dpf::bit::one));
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    static_assert(P0::is_multilevel, "at<> key must be multi-level");
    ASSERT_GT(P0::num_outputs, 0u);

    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));
    EXPECT_TRUE(static_cast<bool>(dpf::reconstruct(
        *dpf::eval_point(dpf::out<0, 10>, r0, x),
        *dpf::eval_point(dpf::out<0, 10>, r1, x))));
    const uint32_t neighbor = x ^ (1u << (32 - 10));
    EXPECT_FALSE(static_cast<bool>(dpf::reconstruct(
        *dpf::eval_point(dpf::out<0, 10>, r0, neighbor),
        *dpf::eval_point(dpf::out<0, 10>, r1, neighbor))));
}

// Point output and comparison channel on one key.
TEST(IncrementalJsonTest, OutputAndCmpRoundTrips)
{
    const uint32_t alpha = 0x00abcdefu;
    const uint64_t yt = 42u;
    auto [k0, k1] = dpf::make_dpf(alpha, uint32_t{7}, dpf::lt(yt));
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));

    EXPECT_EQ(static_cast<uint32_t>(dpf::reconstruct(
        *dpf::eval_point(r0, alpha), *dpf::eval_point(r1, alpha))), 7u);
    EXPECT_EQ(static_cast<uint32_t>(dpf::reconstruct(
        *dpf::eval_point(r0, alpha ^ 1u), *dpf::eval_point(r1, alpha ^ 1u))), 0u);

    const uint64_t mask = k0.cmp().mask;
    auto recon_cmp = [&](uint32_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, r0, q),
            dpf::eval_point(dpf::cmp, r1, q)) & mask;
    };
    EXPECT_EQ(recon_cmp(alpha - 1u), yt);
    EXPECT_EQ(recon_cmp(alpha), 0u);
    EXPECT_EQ(recon_cmp(alpha + 1u), 0u);
}

// Comparison payloads wider than 64 bits stay in the correction words.
TEST(IncrementalJsonTest, Uint128CmpRoundTrips)
{
    using beta = simde_uint128;
    const std::uint8_t alpha = 0x20;
    const beta hi = (beta{1} << 80) + 9;
    const beta lo = beta{3};
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(hi, lo));
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    const std::string s0 = dpf::json::to_json(k0);
    EXPECT_NE(s0.find("\"cw_last\":["), std::string::npos);

    auto r0 = dpf::json::from_json<P0>(s0);
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));
    EXPECT_EQ(r0.cw_last_word(), k0.cw_last_word());
    EXPECT_EQ(r0.cmp_addend_word(), k0.cmp_addend_word());
    EXPECT_EQ(r0.value_cw(), k0.value_cw());
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const beta got = dpf::reconstruct(
            dpf::eval_point<beta>(dpf::cmp, r0, q),
            dpf::eval_point<beta>(dpf::cmp, r1, q));
        const beta want = q > alpha ? hi : lo;
        EXPECT_EQ(got, want) << int(q);
    }
}

// iDCF prefix corrections are part of the same wide word.
TEST(IncrementalJsonTest, Uint128IdcfRoundTrips)
{
    using beta = simde_uint128;
    const std::uint8_t alpha = 0x6e;
    const beta y = (beta{1} << 100) + 13;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::idcf(dpf::gt(y)));
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));
    EXPECT_EQ(r0.prefix_cws(), k0.prefix_cws());
    for (std::uint8_t q : {std::uint8_t{0}, std::uint8_t{0x60}, alpha, std::uint8_t{0x70}})
    {
        const beta got = dpf::reconstruct(
            dpf::eval_point<beta>(dpf::cmp, r0, q),
            dpf::eval_point<beta>(dpf::cmp, r1, q));
        const beta want = dpf::reconstruct(
            dpf::eval_point<beta>(dpf::cmp, k0, q),
            dpf::eval_point<beta>(dpf::cmp, k1, q));
        EXPECT_EQ(got, want) << int(q);
    }
}

// Wildcard coefficients survive, so assign_cmp still works after a round-trip.
TEST(IncrementalJsonTest, WildcardCmpAssignAfterRoundTrip)
{
    const uint32_t alpha = 0x00abcdefu;
    const uint64_t yt = 42u;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lt(dpf::wildcard_value<uint64_t>{}));
    using P0 = std::decay_t<decltype(k0)>;
    using P1 = std::decay_t<decltype(k1)>;
    static_assert(P0::cmp_is_wildcard, "expected a wildcard comparison key");

    auto r0 = dpf::json::from_json<P0>(dpf::json::to_json(k0));
    auto r1 = dpf::json::from_json<P1>(dpf::json::to_json(k1));
    EXPECT_FALSE(r0.cmp_assigned());
    EXPECT_EQ(r0.value_cw_coeff(), k0.value_cw_coeff());
    EXPECT_EQ(r0.cw_last_coeff_word(), k0.cw_last_coeff_word());

    dpf::assign_cmp(r0, r1, yt);
    EXPECT_TRUE(r0.cmp_assigned());
    const uint64_t mask = k0.cmp().mask;
    auto r = [&](uint32_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, r0, q),
            dpf::eval_point(dpf::cmp, r1, q)) & mask;
    };
    EXPECT_EQ(r(alpha - 1u), yt);
    EXPECT_EQ(r(0u), yt);
    EXPECT_EQ(r(alpha), 0u);
    EXPECT_EQ(r(alpha + 1u), 0u);

    auto a0 = dpf::json::from_json<P0>(dpf::json::to_json(r0));
    auto a1 = dpf::json::from_json<P1>(dpf::json::to_json(r1));
    EXPECT_TRUE(a0.cmp_assigned());
    auto again = [&](uint32_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, a0, q),
            dpf::eval_point(dpf::cmp, a1, q)) & mask;
    };
    EXPECT_EQ(again(alpha - 1u), yt);
    EXPECT_EQ(again(alpha), 0u);
}

}  // namespace