#include <gtest/gtest.h>

#include <cstdint>
#include <stdexcept>
#include <vector>

#include "dpf.hpp"

namespace
{

auto make_unit16(std::uint16_t alpha)
{
    return dpf::make_dpf(alpha, dpf::idpf_ones<16>());
}

}  // namespace

TEST(EvalUntil, SpineMatchesEvalPoint)
{
    const std::uint16_t alpha = 0xBEEF;
    auto [k0, k1] = make_unit16(alpha);

    dpf::idpf_eval_ctx ctx0(k0);
    dpf::idpf_eval_ctx ctx1(k1);
    EXPECT_EQ(ctx0.level(), 0u);
    EXPECT_EQ(ctx0.node_count(), 1u);

    // Empty prefixes on a fresh context leave the root in place.
    auto empty = dpf::eval_until(ctx0, 0, std::vector<std::uint16_t>{});
    EXPECT_TRUE(empty.empty());
    EXPECT_EQ(ctx0.node_count(), 1u);

    std::uint16_t prefix = 0;
    for (std::size_t level = 1; level <= 16; ++level)
    {
        const auto bit = static_cast<std::uint16_t>(
            (alpha >> (16 - level)) & 1u);
        prefix = static_cast<std::uint16_t>((prefix << 1) | bit);
        auto s0 = dpf::eval_until(ctx0, level,
            std::vector<std::uint16_t>{prefix});
        auto s1 = dpf::eval_until(ctx1, level,
            std::vector<std::uint16_t>{prefix});
        ASSERT_EQ(s0.size(), 1u);
        ASSERT_EQ(s1.size(), 1u);
        EXPECT_EQ(ctx0.node_count(), 1u);
        EXPECT_LE(ctx0.node_count(), level);

        const auto domain = static_cast<std::uint16_t>(
            prefix << (16 - level));
        // Dispatch eval_point through a level switch for the matching out<I,N>.
        std::uint64_t ref = 0;
        switch (level)
        {
#define LIBDPF_CHECK_LEVEL(L)                                                  \
        case L:                                                                \
            ref = dpf::reconstruct(                                            \
                *dpf::eval_point(dpf::out<L - 1, L>, k0, domain),              \
                *dpf::eval_point(dpf::out<L - 1, L>, k1, domain));             \
            break
        LIBDPF_CHECK_LEVEL(1); LIBDPF_CHECK_LEVEL(2); LIBDPF_CHECK_LEVEL(3);
        LIBDPF_CHECK_LEVEL(4); LIBDPF_CHECK_LEVEL(5); LIBDPF_CHECK_LEVEL(6);
        LIBDPF_CHECK_LEVEL(7); LIBDPF_CHECK_LEVEL(8); LIBDPF_CHECK_LEVEL(9);
        LIBDPF_CHECK_LEVEL(10); LIBDPF_CHECK_LEVEL(11); LIBDPF_CHECK_LEVEL(12);
        LIBDPF_CHECK_LEVEL(13); LIBDPF_CHECK_LEVEL(14); LIBDPF_CHECK_LEVEL(15);
        LIBDPF_CHECK_LEVEL(16);
#undef LIBDPF_CHECK_LEVEL
        default:
            FAIL() << "bad level";
        }
        EXPECT_EQ(dpf::reconstruct(s0[0], s1[0]), ref);
        EXPECT_EQ(ref, 1u);
    }
}

TEST(EvalUntil, RejectsNonExtendingPrefix)
{
    auto [k0, k1] = make_unit16(0x0001);
    (void)k1;
    dpf::idpf_eval_ctx ctx(k0);
    dpf::eval_until(ctx, 1, std::vector<std::uint16_t>{0});
    EXPECT_THROW(
        dpf::eval_until(ctx, 2, std::vector<std::uint16_t>{3}),
        std::invalid_argument);
}

TEST(EvalUntil, BothChildrenStayLinear)
{
    auto [k0, k1] = make_unit16(0xF00D);
    dpf::idpf_eval_ctx ctx0(k0);
    dpf::idpf_eval_ctx ctx1(k1);
    auto s0 = dpf::eval_until(ctx0, 1, std::vector<std::uint16_t>{0, 1});
    auto s1 = dpf::eval_until(ctx1, 1, std::vector<std::uint16_t>{0, 1});
    EXPECT_EQ(ctx0.node_count(), 2u);
    EXPECT_EQ(dpf::reconstruct(s0[0], s1[0]), 0u);
    EXPECT_EQ(dpf::reconstruct(s0[1], s1[1]), 1u);
}

// Seam with the older full-prefix walk: at depth 1, eval_until on {0,1}
// opens the same counts as eval_prefixes(out<0,1>).
TEST(EvalUntil, MatchesEvalPrefixesAtDepthOne)
{
    const std::uint16_t alpha = 0x8000;
    auto [k0, k1] = make_unit16(alpha);

    dpf::idpf_eval_ctx ctx0(k0);
    dpf::idpf_eval_ctx ctx1(k1);
    auto u0 = dpf::eval_until(ctx0, 1, std::vector<std::uint16_t>{0, 1});
    auto u1 = dpf::eval_until(ctx1, 1, std::vector<std::uint16_t>{0, 1});

    auto [b0, it0] = dpf::eval_prefixes(dpf::out<0, 1>, k0);
    auto [b1, it1] = dpf::eval_prefixes(dpf::out<0, 1>, k1);
    (void)b0;
    (void)b1;
    auto p0 = std::begin(it0);
    auto p1 = std::begin(it1);
    EXPECT_EQ(dpf::reconstruct(u0[0], u1[0]), dpf::reconstruct(*p0, *p1));
    ++p0;
    ++p1;
    EXPECT_EQ(dpf::reconstruct(u0[1], u1[1]), dpf::reconstruct(*p0, *p1));
}

TEST(EvalUntil, RetainThenContinue)
{
    auto [k0, k1] = make_unit16(0xC000);
    dpf::idpf_eval_ctx ctx0(k0);
    dpf::idpf_eval_ctx ctx1(k1);
    dpf::eval_until(ctx0, 1, std::vector<std::uint16_t>{0, 1});
    dpf::eval_until(ctx1, 1, std::vector<std::uint16_t>{0, 1});
    ctx0.retain(1);
    ctx1.retain(1);
    EXPECT_EQ(ctx0.node_count(), 1u);
    auto s0 = dpf::eval_until(ctx0, 2, std::vector<std::uint16_t>{2, 3});
    auto s1 = dpf::eval_until(ctx1, 2, std::vector<std::uint16_t>{2, 3});
    // 0xC000 = 1100… so length-2 prefix is 3 (bits 11).
    EXPECT_EQ(dpf::reconstruct(s0[0], s1[0]), 0u);
    EXPECT_EQ(dpf::reconstruct(s0[1], s1[1]), 1u);
}