// Regression checks for the eval / memoizer / output-buffer audit fixes.
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <limits>
#include <vector>

#include "dpf.hpp"

static int fails = 0;

static void expect(bool ok, const char *what)
{
    if (!ok)
    {
        std::fprintf(stderr, "FAIL: %s\n", what);
        ++fails;
    }
}

static void test_wrap_interval()
{
    using input_t = uint8_t;
    using output_t = uint64_t;
    const input_t alpha = 252;
    const output_t beta = 0x1234567890abcdefULL;
    auto [k0, k1] = dpf::make_dpf(input_t{alpha}, output_t{beta});

    const input_t from = 250;
    const input_t to = 5;
    auto [buf0, it0] = dpf::eval_interval(k0, from, to);
    auto [buf1, it1] = dpf::eval_interval(k1, from, to);
    auto z0 = std::begin(it0);
    auto z1 = std::begin(it1);
    int seen = 0;
    for (int step = 0; step < 16; ++step, ++z0, ++z1)
    {
        const input_t x = static_cast<input_t>(from + step);
        if (z0 == std::end(it0))
            break;
        const auto p0 = dpf::eval_point(k0, x);
        const auto p1 = dpf::eval_point(k1, x);
        const output_t interval = static_cast<output_t>(*z1) - static_cast<output_t>(*z0);
        const output_t point = static_cast<output_t>(*p1) - static_cast<output_t>(*p0);
        const output_t want = (x == alpha) ? beta : output_t{0};
        expect(interval == point && interval == want, "wrap interval matches point");
        ++seen;
    }
    expect(z0 == std::end(it0), "wrap iterable consumed");
    expect(seen == 12, "wrap covers 250..255,0..5");
}

static void test_uint64_suffix()
{
    using input_t = uint64_t;
    using output_t = uint64_t;
    using key_t = dpf::utils::dpf_type_t<dpf::prg::aes128, dpf::prg::aes128,
        input_t, output_t>;
    const input_t alpha = std::numeric_limits<input_t>::max() - 1;
    const output_t beta = 77;
    auto [k0, k1] = dpf::make_dpf(input_t{alpha}, output_t{beta});

    input_t flipped_to = std::numeric_limits<input_t>::max();
    dpf::utils::flip_msb_if_signed_integral(flipped_to);
    const auto to_node = dpf::utils::get_to_node<key_t>(flipped_to);
    expect(to_node == (input_t{1} << 63), "uint64 max exclusive leaf is 2^63");

    const input_t from = std::numeric_limits<input_t>::max() - 3;
    const input_t to = std::numeric_limits<input_t>::max();
    auto [buf0, it0] = dpf::eval_interval(k0, from, to);
    auto [buf1, it1] = dpf::eval_interval(k1, from, to);
    auto z0 = std::begin(it0);
    auto z1 = std::begin(it1);
    int seen = 0;
    for (input_t x = from;; ++x, ++z0, ++z1)
    {
        expect(z0 != std::end(it0), "uint64 suffix still has outputs");
        const auto p0 = dpf::eval_point(k0, x);
        const auto p1 = dpf::eval_point(k1, x);
        const output_t interval = static_cast<output_t>(*z1) - static_cast<output_t>(*z0);
        const output_t point = static_cast<output_t>(*p1) - static_cast<output_t>(*p0);
        const output_t want = (x == alpha) ? beta : output_t{0};
        expect(interval == point && interval == want, "uint64 suffix matches point");
        ++seen;
        if (x == to)
            break;
    }
    ++z0;
    expect(z0 == std::end(it0), "uint64 suffix iterable consumed");
    expect(seen == 4, "uint64 suffix length");
}

static void test_narrow_inner_product()
{
    using input_t = uint8_t;
    using output_t = uint8_t;
    const input_t alpha = 30;
    const output_t beta = 7;
    auto [k0, k1] = dpf::make_dpf(input_t{alpha}, output_t{beta});
    uint64_t w[256];
    for (int i = 0; i < 256; ++i)
        w[i] = static_cast<uint64_t>(i * 3 + 1);

    const input_t from = 0;
    const input_t to = 255;
    auto memo0 = dpf::make_basic_full_memoizer(k0);
    auto memo1 = dpf::make_basic_full_memoizer(k1);
    auto ip0 = dpf::eval_inner_product(k0, from, to, w, memo0);
    auto ip1 = dpf::eval_inner_product(k1, from, to, w, memo1);
    // Shares live in the output group, so a uint8 dot product is mod 256.
    // alpha sits in lane 14 of its leaf, past the old 8-lane read.
    const uint64_t got = static_cast<uint64_t>(ip1) - static_cast<uint64_t>(ip0);
    const uint64_t want = (static_cast<uint64_t>(beta) * w[alpha]) & 0xffu;
    expect(got == want, "uint8 inner product (16 lanes per leaf)");
}

static void test_empty_sequence()
{
    using input_t = uint8_t;
    using output_t = uint64_t;
    auto [k0, k1] = dpf::make_dpf(input_t{3}, output_t{1});
    std::vector<input_t> pts;
    auto [buf, it] = dpf::eval_sequence(k0, pts.begin(), pts.end(),
        dpf::return_output_only_tag_{});
    expect(std::begin(it) == std::end(it), "empty output-only sequence");
    auto [bbuf, bit] = dpf::eval_sequence_breadth_first(k0, pts.begin(), pts.end());
    expect(std::begin(bit) == std::end(bit), "empty breadth-first sequence");
    auto recipe = dpf::make_sequence_recipe<decltype(k0)>(pts.begin(), pts.end());
    expect(recipe.num_leaf_nodes() == 0, "empty recipe has no leaves");
    auto [rbuf, rit] = dpf::eval_sequence(k1, recipe, dpf::return_output_only_tag_{});
    expect(std::begin(rit) == std::end(rit), "empty recipe eval");
    (void)buf;
    (void)bbuf;
    (void)rbuf;
}

static void test_path_high_water()
{
    using input_t = uint16_t;
    using output_t = uint64_t;
    const input_t alpha = 1000;
    const output_t beta = 42;
    auto [k0, k1] = dpf::make_dpf(input_t{alpha}, output_t{beta});
    auto path = dpf::make_basic_path_memoizer(k0);
    input_t tx = alpha;
    dpf::utils::flip_msb_if_signed_integral(tx);
    dpf::detail::ensure_level(k0, tx, path, 2);
    const auto partial = dpf::eval_point(k0, alpha, path);
    const auto fresh = dpf::eval_point(k0, alpha);
    const auto other = dpf::eval_point(k1, alpha);
    const output_t got = static_cast<output_t>(*other) - static_cast<output_t>(*partial);
    const output_t want = static_cast<output_t>(*other) - static_cast<output_t>(*fresh);
    expect(got == beta && want == beta, "point eval resumes a partial path");
}

int main()
{
    test_wrap_interval();
    test_uint64_suffix();
    test_narrow_inner_product();
    test_empty_sequence();
    test_path_high_water();
    if (fails != 0)
    {
        std::fprintf(stderr, "%d check(s) failed\n", fails);
        return EXIT_FAILURE;
    }
    std::printf("eval audit fixes ok\n");
    return EXIT_SUCCESS;
}