// Lane Z: the RNG word ledger.
//
// The whole point of the instrument is that a word delta is exact whatever spent the words, so
// the tests drive a real MT19937 by known amounts and check the ledger's arithmetic against the
// count -- including across block boundaries, across a NextInt rejection loop, and when the
// observations arrive out of chronological order (which is how Hook<> renders nested calls).
#include "shim/hooks/rng_ledger.h"

#include <cstring>
#include <vector>

#include "check.h"
#include "mars/rng/mt19937.h"

using shim::hooks::RngLedger;
using shim::hooks::RngPos;
using mars::rng::MT19937;

namespace {

RngPos observe(RngLedger& L, const MT19937& g) { return L.observe(g.state(), g.left()); }

void test_delta_within_one_block() {
    RngLedger L;
    L.reset();
    MT19937 g(12345);
    const RngPos a = observe(L, g);
    CHECK(a.known);
    for (int i = 0; i < 100; ++i) (void)g.next_u32();
    const RngPos b = observe(L, g);
    CHECK(b.known);
    CHECK_EQ(b.words - a.words, 100LL);
    CHECK_EQ(b.block, a.block);
}

void test_delta_across_blocks() {
    RngLedger L;
    L.reset();
    MT19937 g(999);
    const RngPos a = observe(L, g);
    // Well past a block boundary, and not a multiple of 624 -- an off-by-one in the position
    // formula would survive a multiple.
    const int n = 624 * 3 + 17;
    for (int i = 0; i < n; ++i) (void)g.next_u32();
    const RngPos b = observe(L, g);
    CHECK(b.known);
    CHECK_EQ(b.words - a.words, static_cast<long long>(n));
    CHECK_EQ(b.block - a.block, 3);
}

void test_exhausted_block_boundary() {
    // left == 0 is a real state: the block is spent and the next draw twists. Position must be
    // continuous across it, or every ledger entry straddling a boundary is off by 624.
    RngLedger L;
    L.reset();
    MT19937 g(7);
    const RngPos a = observe(L, g);
    for (int i = 0; i < 624; ++i) (void)g.next_u32();
    const RngPos b = observe(L, g);
    CHECK_EQ(b.left, 0);
    CHECK_EQ(b.block, a.block);  // still the same block; the twist has not happened yet
    CHECK_EQ(b.words - a.words, 624LL);
    (void)g.next_u32();
    const RngPos c = observe(L, g);
    CHECK_EQ(c.block, a.block + 1);
    CHECK_EQ(c.words - a.words, 625LL);
}

void test_rejection_loop_counts_words_not_draws() {
    // next_int_inclusive can spend several words on one call. The ledger must report the words.
    RngLedger L;
    L.reset();
    MT19937 g(4242);
    MT19937 shadow(4242);
    const RngPos a = observe(L, g);
    (void)g.next_int_inclusive(100);
    const RngPos b = observe(L, g);
    long long words = 0;
    for (;;) {
        const std::uint32_t r = shadow.next_u32() & MT19937::cover_mask(100);
        ++words;
        if (r <= 100) break;
    }
    CHECK(words >= 1);
    CHECK_EQ(b.words - a.words, words);
}

void test_out_of_order_observation() {
    // Hook<> renders a nested call's snapshots before the outer call's. The outer `before`
    // state is therefore observed twice: once at entry (in order) and once at render time,
    // by which point the chain has moved on. The second lookup must still resolve.
    RngLedger L;
    L.reset();
    MT19937 g(31337);
    std::uint32_t outer_block[624];
    std::memcpy(outer_block, g.state(), sizeof outer_block);
    const int outer_left = g.left();
    const RngPos entry = L.observe(outer_block, outer_left);  // observed at entry, in order
    CHECK(entry.known);

    for (int i = 0; i < 624 * 2 + 5; ++i) (void)g.next_u32();
    const RngPos inner_after = observe(L, g);  // rendered first
    CHECK(inner_after.known);
    for (int i = 0; i < 30; ++i) (void)g.next_u32();
    const RngPos outer_after = observe(L, g);
    CHECK(outer_after.known);

    // ... and now the stale `before` snapshot is rendered.
    const RngPos rendered = L.observe(outer_block, outer_left);
    CHECK(rendered.known);
    CHECK_EQ(rendered.words, entry.words);
    CHECK_EQ(outer_after.words - rendered.words, static_cast<long long>(624 * 2 + 35));
}

void test_backwards_without_entry_observation_is_unknown() {
    // The honest failure: a state behind the anchor cannot be positioned, and the ledger says
    // so rather than inventing a number. This is why hooks observe at entry.
    RngLedger L;
    L.reset();
    MT19937 g(555);
    std::uint32_t early[624];
    std::memcpy(early, g.state(), sizeof early);
    const int early_left = g.left();
    for (int i = 0; i < 624 * 4; ++i) (void)g.next_u32();
    const RngPos anchor = observe(L, g);  // the chain starts HERE
    CHECK(anchor.known);
    const RngPos behind = L.observe(early, early_left);
    CHECK(!behind.known);
    CHECK(L.misses() >= 1);
}

void test_second_generator_reads_unknown() {
    // One ledger, two independent generators: the second one's blocks are not on the first
    // one's chain. That must read unknown, because it is the signal that the "one strategic
    // generator" assumption failed.
    RngLedger L;
    L.reset();
    MT19937 a(1);
    MT19937 b(2);
    CHECK(observe(L, a).known);
    CHECK(!observe(L, b).known);
}

void test_bad_left_is_rejected() {
    RngLedger L;
    L.reset();
    MT19937 g(8);
    CHECK(!L.observe(g.state(), -1).known);
    CHECK(!L.observe(g.state(), 625).known);
    CHECK(!L.observe(nullptr, 100).known);
}

void test_observe_object_layout() {
    // The live-memory path: vptr, mt[624] at +4, next at +0x9c4, left at +0x9c8 (RNG_size
    // 0x9cc). A short region must be refused rather than read past its end.
    RngLedger L;
    L.reset();
    MT19937 g(2024);
    std::vector<char> obj(0x9cc, 0);
    std::memcpy(obj.data() + 4, g.state(), 624 * 4);
    const std::int32_t left = g.left();
    std::memcpy(obj.data() + 0x9c8, &left, 4);
    const RngPos a = L.observe_object(obj.data(), obj.size());
    CHECK(a.known);
    CHECK_EQ(a.left, left);
    CHECK(!L.observe_object(obj.data(), 0x100).known);
}

}  // namespace

int main() {
    test_delta_within_one_block();
    test_delta_across_blocks();
    test_exhausted_block_boundary();
    test_rejection_loop_counts_words_not_draws();
    test_out_of_order_observation();
    test_backwards_without_entry_observation_is_unknown();
    test_second_generator_reads_unknown();
    test_bad_left_is_rejected();
    test_observe_object_layout();
    return simtest::finish("shim_rng_ledger_unit");
}