#include <paritypp/galois.h>
#include <paritypp/matrix.h>
#include <paritypp/parity.h>
#include <paritypp/auth.h>
#include <paritypp/protocol.h>
#include <paritypp/server.h>
#include <paritypp/client.h>

#include <iostream>
#include <cassert>
#include <vector>
#include <algorithm>
#include <cstring>
#include <random>
#include <filesystem>
#include <thread>
#include <atomic>
#include <cstdlib>
#include <numeric>
#include <bitset>

using namespace paritypp;

// ============================================================================
// Helpers
// ============================================================================

static std::mt19937 rng(42); // deterministic seed for reproducibility

static std::vector<uint8_t> random_data(size_t len) {
    std::vector<uint8_t> buf(len);
    std::uniform_int_distribution<int> dist(0, 255);
    for (auto& b : buf) b = static_cast<uint8_t>(dist(rng));
    return buf;
}

// ============================================================================
// 1. Round-trip consistency: encode -> decode must recover original data
//    for every combination of lost blocks up to m
// ============================================================================

static void test_roundtrip_all_loss_patterns() {
    std::cout << "=== Round-trip: all loss patterns (k=3, m=2) ===" << std::endl;

    galois::init();
    parity p(3, 2);
    size_t block_size = 128;

    std::vector<std::vector<uint8_t>> data(3);
    for (size_t i = 0; i < 3; i++)
        data[i] = random_data(block_size);

    auto par = p.encode(data);
    assert(p.verify(data, par));

    // All blocks: indices 0,1,2 = data; 3,4 = parity
    std::vector<std::vector<uint8_t>> all_blocks;
    for (auto& d : data) all_blocks.push_back(d);
    for (auto& pb : par) all_blocks.push_back(pb);

    size_t n = all_blocks.size(); // 5

    // Try every subset of size >= k (3) as surviving blocks
    for (uint32_t mask = 0; mask < (1u << n); ++mask) {
        size_t popcount = std::bitset<32>(mask).count();
        if (popcount < 3) continue; // need at least k surviving

        std::vector<std::vector<uint8_t>> surviving;
        std::vector<size_t> indices;
        for (size_t i = 0; i < n; ++i) {
            if (mask & (1u << i)) {
                surviving.push_back(all_blocks[i]);
                indices.push_back(i);
            }
        }

        auto decoded = p.decode(surviving, indices, block_size);
        for (size_t i = 0; i < 3; ++i) {
            assert(decoded[i] == data[i]);
        }
    }

    std::cout << "  All loss pattern round-trips passed." << std::endl;
}

// ============================================================================
// 2. Split/join consistency with various data sizes
// ============================================================================

static void test_split_join_edge_cases() {
    std::cout << "=== Split/Join edge cases ===" << std::endl;

    // Single byte
    {
        std::vector<uint8_t> data = {0x42};
        auto blocks = parity::split(data.data(), data.size(), 4);
        assert(blocks.size() == 4);
        auto restored = parity::join(blocks, data.size());
        assert(restored == data);
    }

    // Size not divisible by k
    for (size_t k = 2; k <= 8; ++k) {
        for (size_t len : {1, 7, 13, 100, 255, 1000, 4096}) {
            auto data = random_data(len);
            auto blocks = parity::split(data.data(), data.size(), k);
            assert(blocks.size() == k);
            // All blocks same size
            for (size_t i = 1; i < k; ++i)
                assert(blocks[i].size() == blocks[0].size());
            auto restored = parity::join(blocks, data.size());
            assert(restored == data);
        }
    }

    // All zeros
    {
        std::vector<uint8_t> data(256, 0x00);
        auto blocks = parity::split(data.data(), data.size(), 3);
        auto restored = parity::join(blocks, data.size());
        assert(restored == data);
    }

    // All 0xFF
    {
        std::vector<uint8_t> data(256, 0xFF);
        auto blocks = parity::split(data.data(), data.size(), 3);
        auto restored = parity::join(blocks, data.size());
        assert(restored == data);
    }

    std::cout << "  Split/Join edge cases passed." << std::endl;
}

// ============================================================================
// 3. End-to-end data consistency with various k,m configurations
// ============================================================================

static void test_encode_decode_configurations() {
    std::cout << "=== Encode/Decode across configurations ===" << std::endl;
    galois::init();

    struct config { size_t k; size_t m; };
    std::vector<config> configs = {
        {2, 1}, {2, 2}, {3, 1}, {3, 2}, {3, 3},
        {4, 1}, {4, 2}, {4, 4}, {5, 3}, {8, 4},
        {10, 5}, {16, 8}
    };

    for (auto& cfg : configs) {
        parity p(cfg.k, cfg.m);
        size_t block_size = 64;

        std::vector<std::vector<uint8_t>> data(cfg.k);
        for (size_t i = 0; i < cfg.k; ++i)
            data[i] = random_data(block_size);

        auto par = p.encode(data);
        assert(par.size() == cfg.m);
        assert(p.verify(data, par));

        // Lose exactly m blocks (worst case), first m blocks
        std::vector<std::vector<uint8_t>> surviving;
        std::vector<size_t> indices;
        for (size_t i = cfg.m; i < cfg.k; ++i) {
            surviving.push_back(data[i]);
            indices.push_back(i);
        }
        for (size_t i = 0; i < cfg.m; ++i) {
            surviving.push_back(par[i]);
            indices.push_back(cfg.k + i);
        }

        auto decoded = p.decode(surviving, indices, block_size);
        for (size_t i = 0; i < cfg.k; ++i)
            assert(decoded[i] == data[i]);
    }

    std::cout << "  All k,m configurations passed." << std::endl;
}

// ============================================================================
// 4. Verify detects single-bit corruption in data and parity blocks
// ============================================================================

static void test_verify_detects_corruption() {
    std::cout << "=== Verify detects corruption ===" << std::endl;
    galois::init();

    parity p(4, 2);
    size_t block_size = 256;

    std::vector<std::vector<uint8_t>> data(4);
    for (size_t i = 0; i < 4; ++i)
        data[i] = random_data(block_size);

    auto par = p.encode(data);
    assert(p.verify(data, par));

    // Single-bit flip in each data block
    for (size_t i = 0; i < 4; ++i) {
        auto corrupted = data;
        corrupted[i][0] ^= 0x01;
        assert(!p.verify(corrupted, par));
    }

    // Single-bit flip in each parity block
    for (size_t i = 0; i < 2; ++i) {
        auto corrupted_par = par;
        corrupted_par[i][0] ^= 0x01;
        assert(!p.verify(data, corrupted_par));
    }

    // Flip last bit of last byte
    {
        auto corrupted = data;
        corrupted[3][block_size - 1] ^= 0x80;
        assert(!p.verify(corrupted, par));
    }

    std::cout << "  Corruption detection passed." << std::endl;
}

// ============================================================================
// 5. Full end-to-end: split -> encode -> lose -> decode -> join
//    with various data sizes
// ============================================================================

static void test_end_to_end_various_sizes() {
    std::cout << "=== End-to-end with various data sizes ===" << std::endl;
    galois::init();

    size_t k = 4, m = 2;
    parity p(k, m);

    for (size_t size : {1, 15, 64, 255, 512, 1000, 4096, 8192}) {
        auto original = random_data(size);
        auto blocks = parity::split(original.data(), original.size(), k);
        auto par = p.encode(blocks);
        assert(p.verify(blocks, par));

        // Lose first 2 data blocks
        std::vector<std::vector<uint8_t>> surviving = {
            blocks[2], blocks[3], par[0], par[1]
        };
        std::vector<size_t> indices = {2, 3, 4, 5};

        auto decoded = p.decode(surviving, indices, blocks[0].size());
        auto result = parity::join(decoded, original.size());
        assert(result == original);
    }

    std::cout << "  End-to-end various sizes passed." << std::endl;
}

// ============================================================================
// 6. Protocol encode/decode round-trip consistency
// ============================================================================

static void test_protocol_roundtrip() {
    std::cout << "=== Protocol encode/decode round-trip ===" << std::endl;

    // store_block / parse_store_block
    {
        auto data = random_data(512);
        uint64_t gid = 0xDEADBEEF12345678ULL;
        uint32_t bidx = 42;

        auto payload = protocol::make_store_block(gid, bidx, data.data(), data.size());

        uint64_t p_gid;
        uint32_t p_bidx;
        const uint8_t* p_data;
        size_t p_len;
        assert(protocol::parse_store_block(payload.data(), payload.size(),
                                            p_gid, p_bidx, p_data, p_len));
        assert(p_gid == gid);
        assert(p_bidx == bidx);
        assert(p_len == data.size());
        assert(memcmp(p_data, data.data(), data.size()) == 0);
    }

    // fetch_block / parse_fetch_block
    {
        uint64_t gid = 0x1234567890ABCDEFULL;
        uint32_t bidx = 99;
        auto payload = protocol::make_fetch_block(gid, bidx);

        uint64_t p_gid;
        uint32_t p_bidx;
        assert(protocol::parse_fetch_block(payload.data(), payload.size(),
                                            p_gid, p_bidx));
        assert(p_gid == gid);
        assert(p_bidx == bidx);
    }

    // block_data / parse_block_data
    {
        auto data = random_data(1024);
        uint64_t gid = 777;
        uint32_t bidx = 3;

        auto payload = protocol::make_block_data(gid, bidx, data.data(), data.size());

        uint64_t p_gid;
        uint32_t p_bidx;
        const uint8_t* p_data;
        size_t p_len;
        assert(protocol::parse_block_data(payload.data(), payload.size(),
                                           p_gid, p_bidx, p_data, p_len));
        assert(p_gid == gid);
        assert(p_bidx == bidx);
        assert(p_len == data.size());
        assert(memcmp(p_data, data.data(), data.size()) == 0);
    }

    // status_resp round-trip
    {
        status_info info;
        info.status = node_status::DEGRADED;
        info.block_count = 12345;
        info.group_count = 678;
        info.peers_total = 5;
        info.peers_online = 3;
        info.scrub_running = true;

        auto payload = protocol::make_status_resp(info);
        status_info parsed;
        assert(protocol::parse_status_resp(payload.data(), payload.size(), parsed));
        assert(parsed.status == info.status);
        assert(parsed.block_count == info.block_count);
        assert(parsed.group_count == info.group_count);
        assert(parsed.peers_total == info.peers_total);
        assert(parsed.peers_online == info.peers_online);
        assert(parsed.scrub_running == info.scrub_running);
    }

    // group_list round-trip
    {
        std::vector<uint64_t> groups = {1, 100, 999999, 0xFFFFFFFFFFFFFFFFULL, 0};
        auto payload = protocol::make_group_list(groups);
        std::vector<uint64_t> parsed;
        assert(protocol::parse_group_list(payload.data(), payload.size(), parsed));
        assert(parsed == groups);
    }

    // u64 encoding round-trip
    {
        uint8_t buf[8];
        for (uint64_t val : {0ULL, 1ULL, 255ULL, 65535ULL,
                             0xDEADBEEFULL, 0xFFFFFFFFFFFFFFFFULL}) {
            protocol::encode_u64(buf, val);
            assert(protocol::decode_u64(buf) == val);
        }
    }

    // u32 encoding round-trip
    {
        uint8_t buf[4];
        for (uint32_t val : {0u, 1u, 255u, 65535u, 0xDEADBEEFu, 0xFFFFFFFFu}) {
            protocol::encode_u32(buf, val);
            assert(protocol::decode_u32(buf) == val);
        }
    }

    // Full message envelope round-trip
    {
        auto data = random_data(300);
        auto msg = protocol::encode_message(msg_type::BLOCK_DATA, data);
        msg_type type;
        uint32_t plen;
        assert(protocol::decode_header(msg.data(), msg.size(), type, plen));
        assert(type == msg_type::BLOCK_DATA);
        assert(plen == data.size());
        assert(memcmp(msg.data() + protocol::HEADER_SIZE, data.data(), data.size()) == 0);
    }

    std::cout << "  Protocol round-trip tests passed." << std::endl;
}

// ============================================================================
// 7. memory_block_store consistency
// ============================================================================

static void test_memory_block_store_consistency() {
    std::cout << "=== Memory block store consistency ===" << std::endl;

    memory_block_store store;

    // Store and retrieve many groups
    const size_t num_groups = 50;
    const size_t blocks_per_group = 6;

    std::vector<std::vector<std::vector<uint8_t>>> expected(num_groups);

    for (size_t g = 0; g < num_groups; ++g) {
        expected[g].resize(blocks_per_group);
        for (size_t b = 0; b < blocks_per_group; ++b) {
            auto data = random_data(128);
            expected[g][b] = data;
            assert(store.store(g + 1000, b, data.data(), data.size()));
        }
    }

    // Verify all data round-trips
    for (size_t g = 0; g < num_groups; ++g) {
        auto block_list = store.list_blocks(g + 1000);
        assert(block_list.size() == blocks_per_group);
        for (size_t b = 0; b < blocks_per_group; ++b) {
            std::vector<uint8_t> out;
            assert(store.fetch(g + 1000, b, out));
            assert(out == expected[g][b]);
        }
    }

    assert(store.total_groups() == num_groups);
    assert(store.total_blocks() == num_groups * blocks_per_group);

    // Overwrite and verify
    for (size_t g = 0; g < num_groups; ++g) {
        auto data = random_data(64);
        expected[g][0] = data;
        assert(store.store(g + 1000, 0, data.data(), data.size()));
    }

    for (size_t g = 0; g < num_groups; ++g) {
        std::vector<uint8_t> out;
        assert(store.fetch(g + 1000, 0, out));
        assert(out == expected[g][0]);
    }

    // Fetch non-existent
    {
        std::vector<uint8_t> out;
        assert(!store.fetch(99999, 0, out));
    }

    // Remove group and verify
    assert(store.remove_group(1000));
    {
        auto blocks = store.list_blocks(1000);
        assert(blocks.empty());
        std::vector<uint8_t> out;
        assert(!store.fetch(1000, 0, out));
    }

    // list_groups
    auto groups = store.list_groups();
    assert(std::find(groups.begin(), groups.end(), 1000) == groups.end());
    assert(groups.size() == num_groups - 1);

    std::cout << "  Memory block store consistency passed." << std::endl;
}

// ============================================================================
// 8. memory_block_store thread safety
// ============================================================================

static void test_memory_block_store_threaded() {
    std::cout << "=== Memory block store thread safety ===" << std::endl;

    memory_block_store store;
    const size_t num_threads = 8;
    const size_t ops_per_thread = 200;
    std::atomic<int> errors{0};

    std::vector<std::thread> threads;
    for (size_t t = 0; t < num_threads; ++t) {
        threads.emplace_back([&, t]() {
            std::mt19937 local_rng(t * 1000);
            std::uniform_int_distribution<int> dist(0, 255);

            for (size_t i = 0; i < ops_per_thread; ++i) {
                uint64_t gid = t * 10000 + i;
                std::vector<uint8_t> data(64);
                for (auto& b : data) b = static_cast<uint8_t>(dist(local_rng));

                if (!store.store(gid, 0, data.data(), data.size())) {
                    errors++;
                    continue;
                }

                std::vector<uint8_t> out;
                if (!store.fetch(gid, 0, out)) {
                    errors++;
                    continue;
                }
                if (out != data) {
                    errors++;
                }
            }
        });
    }

    for (auto& t : threads) t.join();
    assert(errors == 0);

    std::cout << "  Memory block store thread safety passed." << std::endl;
}

// ============================================================================
// 9. file_block_store data consistency
// ============================================================================

static void test_file_block_store_consistency() {
    std::cout << "=== File block store data consistency ===" << std::endl;

    std::string dir = "/tmp/paritypp_consistency_" + std::to_string(getpid());
    std::filesystem::remove_all(dir);

    const size_t num_groups = 20;
    const size_t blocks_per_group = 6;

    std::vector<std::vector<std::vector<uint8_t>>> expected(num_groups);

    // Store data
    {
        file_block_store store(dir, 0);
        for (size_t g = 0; g < num_groups; ++g) {
            expected[g].resize(blocks_per_group);
            for (size_t b = 0; b < blocks_per_group; ++b) {
                auto data = random_data(256);
                expected[g][b] = data;
                assert(store.store(g + 500, b, data.data(), data.size()));
            }
        }
    }

    // Reopen and verify all data persisted correctly
    {
        file_block_store store(dir, 0);
        assert(store.total_groups() == num_groups);
        assert(store.total_blocks() == num_groups * blocks_per_group);

        for (size_t g = 0; g < num_groups; ++g) {
            for (size_t b = 0; b < blocks_per_group; ++b) {
                std::vector<uint8_t> out;
                assert(store.fetch(g + 500, b, out));
                assert(out == expected[g][b]);
            }
        }
    }

    // Overwrite some blocks, vacuum, and verify
    {
        file_block_store store(dir, 0);
        for (size_t g = 0; g < num_groups; ++g) {
            auto data = random_data(128);
            expected[g][0] = data;
            assert(store.store(g + 500, 0, data.data(), data.size()));
        }
        assert(store.vacuum());

        for (size_t g = 0; g < num_groups; ++g) {
            for (size_t b = 0; b < blocks_per_group; ++b) {
                std::vector<uint8_t> out;
                assert(store.fetch(g + 500, b, out));
                assert(out == expected[g][b]);
            }
        }
    }

    // Reopen after vacuum and verify again
    {
        file_block_store store(dir, 0);
        for (size_t g = 0; g < num_groups; ++g) {
            for (size_t b = 0; b < blocks_per_group; ++b) {
                std::vector<uint8_t> out;
                assert(store.fetch(g + 500, b, out));
                assert(out == expected[g][b]);
            }
        }
    }

    std::filesystem::remove_all(dir);
    std::cout << "  File block store data consistency passed." << std::endl;
}

// ============================================================================
// 10. file_block_store: buffered writes consistency
// ============================================================================

static void test_file_block_store_buffered() {
    std::cout << "=== File block store buffered writes ===" << std::endl;

    std::string dir = "/tmp/paritypp_buffered_" + std::to_string(getpid());
    std::filesystem::remove_all(dir);

    std::vector<std::vector<uint8_t>> expected(10);

    {
        // Use default buffer threshold (256KB) so writes are coalesced
        file_block_store store(dir);
        for (size_t i = 0; i < 10; ++i) {
            expected[i] = random_data(512);
            assert(store.store(1, i, expected[i].data(), expected[i].size()));
        }

        // Verify reads work even from write buffer
        for (size_t i = 0; i < 10; ++i) {
            std::vector<uint8_t> out;
            assert(store.fetch(1, i, out));
            assert(out == expected[i]);
        }

        store.flush();
    }

    // Reopen and verify all persisted
    {
        file_block_store store(dir, 0);
        for (size_t i = 0; i < 10; ++i) {
            std::vector<uint8_t> out;
            assert(store.fetch(1, i, out));
            assert(out == expected[i]);
        }
    }

    std::filesystem::remove_all(dir);
    std::cout << "  File block store buffered writes passed." << std::endl;
}

// ============================================================================
// 11. file_block_store: remove group then vacuum consistency
// ============================================================================

static void test_file_block_store_remove_vacuum() {
    std::cout << "=== File block store remove + vacuum ===" << std::endl;

    std::string dir = "/tmp/paritypp_rmvac_" + std::to_string(getpid());
    std::filesystem::remove_all(dir);

    {
        file_block_store store(dir, 0);

        // Store 3 groups
        for (uint64_t g = 1; g <= 3; ++g) {
            for (uint32_t b = 0; b < 4; ++b) {
                std::vector<uint8_t> d(64, static_cast<uint8_t>(g * 10 + b));
                store.store(g, b, d.data(), d.size());
            }
        }

        // Remove middle group
        assert(store.remove_group(2));

        // Vacuum to compact
        assert(store.vacuum());

        // Verify group 1 and 3 still intact
        for (uint64_t g : {1ULL, 3ULL}) {
            for (uint32_t b = 0; b < 4; ++b) {
                std::vector<uint8_t> out;
                assert(store.fetch(g, b, out));
                assert(out.size() == 64);
                assert(out[0] == static_cast<uint8_t>(g * 10 + b));
            }
        }

        // Group 2 gone
        assert(store.list_blocks(2).empty());
    }

    // Reopen and recheck
    {
        file_block_store store(dir, 0);
        assert(store.total_groups() == 2);
        for (uint64_t g : {1ULL, 3ULL}) {
            auto blocks = store.list_blocks(g);
            assert(blocks.size() == 4);
        }
        assert(store.list_blocks(2).empty());
    }

    std::filesystem::remove_all(dir);
    std::cout << "  File block store remove + vacuum passed." << std::endl;
}

// ============================================================================
// 12. Galois field: region operations consistency
// ============================================================================

static void test_galois_region_ops() {
    std::cout << "=== Galois region operations ===" << std::endl;
    galois::init();

    size_t len = 256;
    auto src = random_data(len);
    uint8_t coeff = 0x53;

    // mul_region: dst = src * coeff (element-wise)
    {
        std::vector<uint8_t> dst(len, 0);
        galois::mul_region(dst.data(), src.data(), coeff, len);
        for (size_t i = 0; i < len; ++i)
            assert(dst[i] == galois::mul(src[i], coeff));
    }

    // addmul_region: dst += src * coeff
    {
        auto dst = random_data(len);
        auto original_dst = dst;
        galois::addmul_region(dst.data(), src.data(), coeff, len);
        for (size_t i = 0; i < len; ++i) {
            uint8_t expected = galois::add(original_dst[i], galois::mul(src[i], coeff));
            assert(dst[i] == expected);
        }
    }

    // mul_region with coeff=0 -> all zeros
    {
        std::vector<uint8_t> dst(len, 0xFF);
        galois::mul_region(dst.data(), src.data(), 0, len);
        for (size_t i = 0; i < len; ++i)
            assert(dst[i] == 0);
    }

    // mul_region with coeff=1 -> copy
    {
        std::vector<uint8_t> dst(len, 0);
        galois::mul_region(dst.data(), src.data(), 1, len);
        assert(dst == src);
    }

    std::cout << "  Galois region operations passed." << std::endl;
}

// ============================================================================
// 13. Matrix: submatrix and row_slice consistency
// ============================================================================

static void test_matrix_submatrix() {
    std::cout << "=== Matrix submatrix/row_slice ===" << std::endl;
    galois::init();

    auto m = matrix::cauchy(6, 4);

    // row_slice: extract rows 2..4
    auto slice = m.row_slice(2, 3);
    assert(slice.rows() == 3);
    assert(slice.cols() == 4);
    for (size_t r = 0; r < 3; ++r)
        for (size_t c = 0; c < 4; ++c)
            assert(slice.at(r, c) == m.at(r + 2, c));

    // submatrix
    std::vector<size_t> rows = {0, 3, 5};
    std::vector<size_t> cols = {1, 3};
    auto sub = m.submatrix(rows, cols);
    assert(sub.rows() == 3);
    assert(sub.cols() == 2);
    for (size_t r = 0; r < 3; ++r)
        for (size_t c = 0; c < 2; ++c)
            assert(sub.at(r, c) == m.at(rows[r], cols[c]));

    // augment: [I | M] has correct dimensions
    auto id = matrix::identity(6);
    auto augmented = id.augment(m);
    assert(augmented.rows() == 6);
    assert(augmented.cols() == 10);
    // Left half is identity
    for (size_t r = 0; r < 6; ++r)
        for (size_t c = 0; c < 6; ++c)
            assert(augmented.at(r, c) == (r == c ? 1 : 0));
    // Right half is original matrix
    for (size_t r = 0; r < 6; ++r)
        for (size_t c = 0; c < 4; ++c)
            assert(augmented.at(r, c + 6) == m.at(r, c));

    std::cout << "  Matrix submatrix/row_slice passed." << std::endl;
}

// ============================================================================
// 14. Matrix: copy constructor and assignment operator
// ============================================================================

static void test_matrix_copy() {
    std::cout << "=== Matrix copy semantics ===" << std::endl;
    galois::init();

    auto orig = matrix::cauchy(4, 4);

    // Copy constructor
    matrix copy1(orig);
    assert(copy1.rows() == orig.rows());
    assert(copy1.cols() == orig.cols());
    for (size_t r = 0; r < orig.rows(); ++r)
        for (size_t c = 0; c < orig.cols(); ++c)
            assert(copy1.at(r, c) == orig.at(r, c));

    // Assignment operator
    matrix copy2(1, 1);
    copy2 = orig;
    assert(copy2.rows() == orig.rows());
    for (size_t r = 0; r < orig.rows(); ++r)
        for (size_t c = 0; c < orig.cols(); ++c)
            assert(copy2.at(r, c) == orig.at(r, c));

    // Mutating copy doesn't affect original
    copy1.at(0, 0) = 0xFF;
    assert(orig.at(0, 0) != 0xFF || copy1.at(0, 0) == orig.at(0, 0));

    std::cout << "  Matrix copy semantics passed." << std::endl;
}

// ============================================================================
// 15. Parity: encode is deterministic
// ============================================================================

static void test_encode_deterministic() {
    std::cout << "=== Encode determinism ===" << std::endl;
    galois::init();

    parity p(4, 2);
    size_t block_size = 128;

    std::vector<std::vector<uint8_t>> data(4);
    for (size_t i = 0; i < 4; ++i)
        data[i] = random_data(block_size);

    auto par1 = p.encode(data);
    auto par2 = p.encode(data);
    assert(par1 == par2);

    // Different parity instance with same k,m
    parity p2(4, 2);
    auto par3 = p2.encode(data);
    assert(par1 == par3);

    std::cout << "  Encode determinism passed." << std::endl;
}

// ============================================================================
// 16. Auth: HMAC consistency
// ============================================================================

static void test_auth_consistency() {
    std::cout << "=== Auth HMAC consistency ===" << std::endl;

    std::vector<uint8_t> key = random_data(32);
    auto nonce = auth::generate_nonce();

    // Same key + nonce -> same HMAC
    auto hmac1 = auth::compute_hmac(key, nonce);
    auto hmac2 = auth::compute_hmac(key, nonce);
    assert(hmac1 == hmac2);
    assert(auth::verify_hmac(key, nonce, hmac1));

    // Different key -> different HMAC
    auto key2 = random_data(32);
    auto hmac3 = auth::compute_hmac(key2, nonce);
    assert(hmac1 != hmac3);
    assert(!auth::verify_hmac(key, nonce, hmac3));

    // Different nonce -> different HMAC
    auto nonce2 = auth::generate_nonce();
    auto hmac4 = auth::compute_hmac(key, nonce2);
    assert(hmac1 != hmac4);
    assert(!auth::verify_hmac(key, nonce, hmac4));

    // Empty key
    {
        std::vector<uint8_t> empty_key;
        auto hmac = auth::compute_hmac(empty_key, nonce);
        assert(auth::verify_hmac(empty_key, nonce, hmac));
        assert(!auth::verify_hmac(key, nonce, hmac));
    }

    // Nonce uniqueness (two generated nonces should differ)
    auto n1 = auth::generate_nonce();
    auto n2 = auth::generate_nonce();
    assert(n1 != n2);

    std::cout << "  Auth HMAC consistency passed." << std::endl;
}

// ============================================================================
// 17. Batch encode/decode round-trip
// ============================================================================

static void test_batch_encode_decode() {
    std::cout << "=== Batch encode/decode ===" << std::endl;

    std::vector<client::batch_item> items;
    for (uint64_t i = 0; i < 10; ++i) {
        client::batch_item item;
        item.key = i * 1000 + 42;
        item.data = random_data(64 + i * 10);
        items.push_back(std::move(item));
    }

    auto encoded = client::encode_batch(items);
    assert(!encoded.empty());

    std::vector<client::batch_item> decoded;
    assert(client::decode_batch(encoded, decoded));
    assert(decoded.size() == items.size());

    for (size_t i = 0; i < items.size(); ++i) {
        assert(decoded[i].key == items[i].key);
        assert(decoded[i].data == items[i].data);
    }

    // Empty batch
    {
        std::vector<client::batch_item> empty;
        auto enc = client::encode_batch(empty);
        std::vector<client::batch_item> dec;
        assert(client::decode_batch(enc, dec));
        assert(dec.empty());
    }

    // Single item batch
    {
        std::vector<client::batch_item> single = {{999, random_data(1)}};
        auto enc = client::encode_batch(single);
        std::vector<client::batch_item> dec;
        assert(client::decode_batch(enc, dec));
        assert(dec.size() == 1);
        assert(dec[0].key == 999);
        assert(dec[0].data == single[0].data);
    }

    std::cout << "  Batch encode/decode passed." << std::endl;
}

// ============================================================================
// 18. Parity recovery: lose parity blocks only (data all survive)
// ============================================================================

static void test_decode_parity_loss_only() {
    std::cout << "=== Decode with parity loss only ===" << std::endl;
    galois::init();

    parity p(4, 3);
    size_t block_size = 64;

    std::vector<std::vector<uint8_t>> data(4);
    for (size_t i = 0; i < 4; ++i)
        data[i] = random_data(block_size);

    auto par = p.encode(data);

    // All data blocks survive, lose all parity blocks
    std::vector<size_t> indices = {0, 1, 2, 3};
    auto decoded = p.decode(data, indices, block_size);
    for (size_t i = 0; i < 4; ++i)
        assert(decoded[i] == data[i]);

    std::cout << "  Parity loss only decode passed." << std::endl;
}

// ============================================================================
// 19. Large block consistency
// ============================================================================

static void test_large_block_consistency() {
    std::cout << "=== Large block consistency ===" << std::endl;
    galois::init();

    parity p(4, 2);
    size_t block_size = 65536; // 64KB blocks

    std::vector<std::vector<uint8_t>> data(4);
    for (size_t i = 0; i < 4; ++i)
        data[i] = random_data(block_size);

    auto par = p.encode(data);
    assert(p.verify(data, par));

    // Lose blocks 1 and 3
    std::vector<std::vector<uint8_t>> surviving = {data[0], data[2], par[0], par[1]};
    std::vector<size_t> indices = {0, 2, 4, 5};
    auto decoded = p.decode(surviving, indices, block_size);

    for (size_t i = 0; i < 4; ++i)
        assert(decoded[i] == data[i]);

    std::cout << "  Large block consistency passed." << std::endl;
}

// ============================================================================
// 20. file_block_store: concurrent read/write
// ============================================================================

static void test_file_block_store_threaded() {
    std::cout << "=== File block store concurrent access ===" << std::endl;

    std::string dir = "/tmp/paritypp_threaded_" + std::to_string(getpid());
    std::filesystem::remove_all(dir);

    {
        file_block_store store(dir, 0);
        const size_t num_threads = 4;
        const size_t ops = 50;
        std::atomic<int> errors{0};

        std::vector<std::thread> threads;
        for (size_t t = 0; t < num_threads; ++t) {
            threads.emplace_back([&, t]() {
                std::mt19937 local_rng(t * 7777);
                std::uniform_int_distribution<int> dist(0, 255);

                for (size_t i = 0; i < ops; ++i) {
                    uint64_t gid = t * 100000 + i;
                    std::vector<uint8_t> data(128);
                    for (auto& b : data) b = static_cast<uint8_t>(dist(local_rng));

                    if (!store.store(gid, 0, data.data(), data.size())) {
                        errors++;
                        continue;
                    }

                    std::vector<uint8_t> out;
                    if (!store.fetch(gid, 0, out)) {
                        errors++;
                        continue;
                    }
                    if (out != data) {
                        errors++;
                    }
                }
            });
        }

        for (auto& t : threads) t.join();
        assert(errors == 0);
    }

    std::filesystem::remove_all(dir);
    std::cout << "  File block store concurrent access passed." << std::endl;
}

// ============================================================================
// main
// ============================================================================

int main() {
    std::cout << "libparitypp Data Consistency Test Suite" << std::endl;
    std::cout << "=======================================" << std::endl << std::endl;

    test_roundtrip_all_loss_patterns();
    test_split_join_edge_cases();
    test_encode_decode_configurations();
    test_verify_detects_corruption();
    test_end_to_end_various_sizes();
    test_protocol_roundtrip();
    test_memory_block_store_consistency();
    test_memory_block_store_threaded();
    test_file_block_store_consistency();
    test_file_block_store_buffered();
    test_file_block_store_remove_vacuum();
    test_galois_region_ops();
    test_matrix_submatrix();
    test_matrix_copy();
    test_encode_deterministic();
    test_auth_consistency();
    test_batch_encode_decode();
    test_decode_parity_loss_only();
    test_large_block_consistency();
    test_file_block_store_threaded();

    std::cout << std::endl << "All consistency tests passed!" << std::endl;
    return 0;
}
