#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 <iostream>
#include <cstring>
#include <cassert>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <filesystem>
#include <cstdlib>

using namespace paritypp;

static void test_galois() {
    std::cout << "=== GF(2^8) Arithmetic ===" << std::endl;
    galois::init();

    // Addition and subtraction are XOR
    assert(galois::add(0x53, 0xCA) == (0x53 ^ 0xCA));
    assert(galois::sub(0x53, 0xCA) == (0x53 ^ 0xCA));

    // Multiplication identity
    assert(galois::mul(1, 0xAB) == 0xAB);
    assert(galois::mul(0xAB, 1) == 0xAB);
    assert(galois::mul(0, 0xAB) == 0);

    // Division identity
    assert(galois::div(0xAB, 1) == 0xAB);
    assert(galois::div(0, 0xAB) == 0);

    // a * inv(a) = 1
    for (int a = 1; a < 256; a++) {
        uint8_t inv_a = galois::inv(static_cast<uint8_t>(a));
        assert(galois::mul(static_cast<uint8_t>(a), inv_a) == 1);
    }

    // a * b / b = a
    for (int a = 1; a < 256; a++) {
        for (int b = 1; b < 256; b++) {
            uint8_t prod = galois::mul(static_cast<uint8_t>(a), static_cast<uint8_t>(b));
            assert(galois::div(prod, static_cast<uint8_t>(b)) == static_cast<uint8_t>(a));
        }
    }

    std::cout << "  All GF(2^8) tests passed." << std::endl;
}

static void test_matrix() {
    std::cout << "=== Matrix Operations ===" << std::endl;
    galois::init();

    // Identity * Identity = Identity
    auto id = matrix::identity(4);
    auto result = id.multiply(id);
    for (size_t r = 0; r < 4; r++) {
        for (size_t c = 0; c < 4; c++) {
            assert(result.at(r, c) == (r == c ? 1 : 0));
        }
    }

    // Matrix inversion: A * A^-1 = I
    auto vand = matrix::vandermonde(4, 4);
    auto inv = vand.invert();
    auto product = vand.multiply(inv);
    for (size_t r = 0; r < 4; r++) {
        for (size_t c = 0; c < 4; c++) {
            uint8_t expected = (r == c) ? 1 : 0;
            assert(product.at(r, c) == expected);
        }
    }

    // Cauchy matrix invertibility
    auto cauchy_m = matrix::cauchy(4, 4);
    auto cauchy_inv = cauchy_m.invert();
    auto cauchy_prod = cauchy_m.multiply(cauchy_inv);
    for (size_t r = 0; r < 4; r++) {
        for (size_t c = 0; c < 4; c++) {
            assert(cauchy_prod.at(r, c) == (r == c ? 1 : 0));
        }
    }

    std::cout << "  All matrix tests passed." << std::endl;
}

static void test_parity_basic() {
    std::cout << "=== Basic Parity (RAID-5 style: k=3, m=1) ===" << std::endl;

    parity p(3, 1);
    assert(p.data_block_count() == 3);
    assert(p.parity_block_count() == 1);
    assert(p.total_block_count() == 4);

    // Create test data blocks
    size_t block_size = 64;
    std::vector<std::vector<uint8_t>> data(3, std::vector<uint8_t>(block_size));
    for (size_t i = 0; i < 3; i++) {
        for (size_t j = 0; j < block_size; j++) {
            data[i][j] = static_cast<uint8_t>((i * 37 + j * 13 + 7) & 0xFF);
        }
    }

    // Encode
    auto parity_blocks = p.encode(data);
    assert(parity_blocks.size() == 1);
    assert(parity_blocks[0].size() == block_size);

    // Verify
    assert(p.verify(data, parity_blocks));

    // Corrupt and verify fails
    auto corrupted = parity_blocks;
    corrupted[0][0] ^= 0xFF;
    assert(!p.verify(data, corrupted));

    // Decode: lose data block 1, reconstruct from remaining
    std::vector<std::vector<uint8_t>> surviving = {data[0], data[2], parity_blocks[0]};
    std::vector<size_t> indices = {0, 2, 3}; // block 0=data0, 2=data2, 3=parity0
    auto decoded = p.decode(surviving, indices, block_size);
    assert(decoded.size() == 3);
    assert(decoded[1] == data[1]);

    std::cout << "  Basic parity tests passed." << std::endl;
}

static void test_parity_raid6() {
    std::cout << "=== RAID-6 style (k=4, m=2) ===" << std::endl;

    parity p(4, 2);
    size_t block_size = 128;
    std::vector<std::vector<uint8_t>> data(4, std::vector<uint8_t>(block_size));
    for (size_t i = 0; i < 4; i++) {
        for (size_t j = 0; j < block_size; j++) {
            data[i][j] = static_cast<uint8_t>((i * 53 + j * 17 + 11) & 0xFF);
        }
    }

    auto parity_blocks = p.encode(data);
    assert(parity_blocks.size() == 2);
    assert(p.verify(data, parity_blocks));

    // Lose 2 data blocks (worst case for RAID-6)
    std::vector<std::vector<uint8_t>> surviving = {data[0], data[3], parity_blocks[0], parity_blocks[1]};
    std::vector<size_t> indices = {0, 3, 4, 5};
    auto decoded = p.decode(surviving, indices, block_size);
    assert(decoded[1] == data[1]);
    assert(decoded[2] == data[2]);

    std::cout << "  RAID-6 parity tests passed." << std::endl;
}

static void test_parity_flexible() {
    std::cout << "=== Flexible parity (k=6, m=4) ===" << std::endl;

    parity p(6, 4);
    size_t block_size = 256;
    std::vector<std::vector<uint8_t>> data(6, std::vector<uint8_t>(block_size));
    for (size_t i = 0; i < 6; i++) {
        for (size_t j = 0; j < block_size; j++) {
            data[i][j] = static_cast<uint8_t>((i * 71 + j * 23 + 3) & 0xFF);
        }
    }

    auto parity_blocks = p.encode(data);
    assert(parity_blocks.size() == 4);
    assert(p.verify(data, parity_blocks));

    // Lose 4 blocks (max tolerable): data blocks 1,3,5 and parity block 0
    std::vector<std::vector<uint8_t>> surviving = {
        data[0], data[2], data[4],
        parity_blocks[1], parity_blocks[2], parity_blocks[3]
    };
    std::vector<size_t> indices = {0, 2, 4, 7, 8, 9};
    auto decoded = p.decode(surviving, indices, block_size);
    for (size_t i = 0; i < 6; i++) {
        assert(decoded[i] == data[i]);
    }

    std::cout << "  Flexible parity tests passed." << std::endl;
}

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

    std::vector<uint8_t> original(1000);
    for (size_t i = 0; i < original.size(); i++) {
        original[i] = static_cast<uint8_t>(i & 0xFF);
    }

    auto blocks = parity::split(original.data(), original.size(), 4);
    assert(blocks.size() == 4);
    // Each block should be ceil(1000/4) = 250 bytes
    assert(blocks[0].size() == 250);

    auto restored = parity::join(blocks, original.size());
    assert(restored == original);

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

static void test_end_to_end() {
    std::cout << "=== End-to-End: split -> encode -> lose blocks -> decode -> join ===" << std::endl;

    // Simulate full workflow
    std::vector<uint8_t> original(2048);
    for (size_t i = 0; i < original.size(); i++) {
        original[i] = static_cast<uint8_t>((i * 97 + 31) & 0xFF);
    }

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

    // Split into k blocks
    auto data_blocks = parity::split(original.data(), original.size(), k);

    // Encode parity
    auto parity_blocks = p.encode(data_blocks);
    assert(p.verify(data_blocks, parity_blocks));

    // Simulate losing data blocks 1 and 3
    std::vector<std::vector<uint8_t>> surviving = {
        data_blocks[0], data_blocks[2],
        parity_blocks[0], parity_blocks[1]
    };
    std::vector<size_t> indices = {0, 2, 4, 5};

    auto reconstructed = p.decode(surviving, indices, data_blocks[0].size());

    // Join and compare
    auto result = parity::join(reconstructed, original.size());
    assert(result == original);

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

static void test_auth_store() {
    std::cout << "=== Auth Store ===" << std::endl;

    auth_store store;

    // Add clients
    store.add_client("node-alpha", "secret-key-alpha");
    std::vector<uint8_t> key2 = {0x01, 0x02, 0x03, 0x04, 0x05};
    store.add_client("node-beta", key2);

    assert(store.has_client("node-alpha"));
    assert(store.has_client("node-beta"));
    assert(!store.has_client("node-gamma"));

    // Lookup
    std::vector<uint8_t> found_key;
    assert(store.lookup("node-alpha", found_key));
    std::string expected_str = "secret-key-alpha";
    assert(found_key == std::vector<uint8_t>(expected_str.begin(), expected_str.end()));

    assert(store.lookup("node-beta", found_key));
    assert(found_key == key2);

    assert(!store.lookup("nonexistent", found_key));

    // List
    auto clients = store.list_clients();
    assert(clients.size() == 2);

    // Remove
    assert(store.remove_client("node-alpha"));
    assert(!store.has_client("node-alpha"));
    assert(!store.remove_client("node-alpha"));

    std::cout << "  Auth store tests passed." << std::endl;
}

static void test_auth_hmac() {
    std::cout << "=== Auth HMAC Challenge-Response ===" << std::endl;

    // Simulate the challenge-response flow
    std::vector<uint8_t> key = {0x41, 0x42, 0x43, 0x44, 0x45, 0x46};

    // Server generates nonce
    auto nonce = auth::generate_nonce();
    assert(nonce.size() == auth::NONCE_SIZE);

    // Client computes HMAC
    auto client_hmac = auth::compute_hmac(key, nonce);
    assert(client_hmac.size() == auth::HMAC_SIZE);

    // Server verifies
    assert(auth::verify_hmac(key, nonce, client_hmac));

    // Wrong key should fail
    std::vector<uint8_t> wrong_key = {0xFF, 0xFE, 0xFD};
    assert(!auth::verify_hmac(wrong_key, nonce, client_hmac));

    // Tampered HMAC should fail
    auto tampered = client_hmac;
    tampered[0] ^= 0xFF;
    assert(!auth::verify_hmac(key, nonce, tampered));

    // Different nonce should fail
    auto other_nonce = auth::generate_nonce();
    assert(!auth::verify_hmac(key, other_nonce, client_hmac));

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

static void test_auth_protocol() {
    std::cout << "=== Auth Protocol Messages ===" << std::endl;

    // Test AUTH_REQUEST encoding/decoding
    std::string name = "my-client";
    auto req = protocol::make_auth_request(name);
    assert(!req.empty());
    std::string parsed_name;
    assert(protocol::parse_auth_request(req.data(), req.size(), parsed_name));
    assert(parsed_name == name);

    // Test AUTH_CHALLENGE encoding/decoding
    auto nonce = auth::generate_nonce();
    auto challenge = protocol::make_auth_challenge(nonce);
    std::vector<uint8_t> parsed_nonce;
    assert(protocol::parse_auth_challenge(challenge.data(), challenge.size(), parsed_nonce));
    assert(parsed_nonce == nonce);

    // Test AUTH_RESPONSE encoding/decoding
    std::vector<uint8_t> key = {0x01, 0x02, 0x03};
    auto hmac = auth::compute_hmac(key, nonce);
    auto resp = protocol::make_auth_response(hmac);
    std::vector<uint8_t> parsed_hmac;
    assert(protocol::parse_auth_response(resp.data(), resp.size(), parsed_hmac));
    assert(parsed_hmac == hmac);

    // Full message encode/decode
    auto msg = protocol::encode_message(msg_type::AUTH_REQUEST, req);
    msg_type type;
    uint32_t plen;
    assert(protocol::decode_header(msg.data(), msg.size(), type, plen));
    assert(type == msg_type::AUTH_REQUEST);
    assert(plen == req.size());

    std::cout << "  Auth protocol tests passed." << std::endl;
}

static uint32_t crc32_reference(const uint8_t* data, size_t len) {
    static uint32_t table[256];
    static bool init = false;
    if (!init) {
        for (uint32_t i = 0; i < 256; ++i) {
            uint32_t c = i;
            for (int j = 0; j < 8; ++j)
                c = (c >> 1) ^ (0xEDB88320 & (-(c & 1)));
            table[i] = c;
        }
        init = true;
    }
    uint32_t crc = 0xFFFFFFFF;
    for (size_t i = 0; i < len; ++i)
        crc = (crc >> 8) ^ table[(crc ^ data[i]) & 0xFF];
    return crc ^ 0xFFFFFFFF;
}

static void test_crc32() {
    std::cout << "=== CRC32 ===" << std::endl;

    // Test 1: known CRC32 of "123456789" is 0xCBF43926
    {
        const uint8_t data[] = "123456789";
        uint32_t got = file_block_store::crc32(data, 9);
        if (got != 0xCBF43926) {
            std::cerr << "  FATAL: CRC32(\"123456789\") = 0x" << std::hex << got
                      << " expected 0xCBF43926" << std::dec << std::endl;
            std::abort();
        }
        std::cout << "  Known value test passed." << std::endl;
    }

    // Test 2: empty data
    {
        uint32_t got = file_block_store::crc32(nullptr, 0);
        if (got != 0x00000000) {
            std::cerr << "  FATAL: CRC32(empty) = 0x" << std::hex << got
                      << " expected 0x00000000" << std::dec << std::endl;
            std::abort();
        }
        std::cout << "  Empty data test passed." << std::endl;
    }

    // Test 3: single byte
    {
        uint8_t b = 0xFF;
        uint32_t got = file_block_store::crc32(&b, 1);
        uint32_t ref = crc32_reference(&b, 1);
        if (got != ref) {
            std::cerr << "  FATAL: CRC32(0xFF) = 0x" << std::hex << got
                      << " expected 0x" << ref << std::dec << std::endl;
            std::abort();
        }
        std::cout << "  Single byte test passed." << std::endl;
    }

    // Test 4: compare table vs PCLMUL at various sizes (especially around 64-byte threshold)
    {
        std::vector<uint8_t> buf(8192);
        // Fill with pseudo-random but deterministic data
        for (size_t i = 0; i < buf.size(); ++i)
            buf[i] = static_cast<uint8_t>((i * 251 + 17) & 0xFF);

        size_t test_sizes[] = {
            1, 7, 15, 16, 31, 32, 48, 63, 64, 65, 80, 100, 127, 128, 129,
            255, 256, 512, 1000, 1024, 2048, 4096, 4097, 8192
        };
        bool ok = true;
        for (size_t sz : test_sizes) {
            uint32_t got = file_block_store::crc32(buf.data(), sz);
            uint32_t ref = crc32_reference(buf.data(), sz);
            if (got != ref) {
                std::cerr << "  CRC32 MISMATCH at size " << sz
                          << ": got 0x" << std::hex << got
                          << " expected 0x" << ref << std::dec << std::endl;
                ok = false;
            }
        }
        if (!ok) { std::cerr << "  FATAL: CRC32 comparison failed" << std::endl; std::abort(); }
        std::cout << "  Table vs PCLMUL comparison passed (" 
                  << (sizeof(test_sizes)/sizeof(test_sizes[0])) << " sizes)." << std::endl;
    }

    // Test 5: unaligned data (PCLMUL alignment handling)
    {
        std::vector<uint8_t> buf(256 + 16);
        for (size_t i = 0; i < buf.size(); ++i)
            buf[i] = static_cast<uint8_t>(i ^ 0xA5);

        bool ok = true;
        for (int offset = 0; offset < 16; ++offset) {
            uint32_t got = file_block_store::crc32(buf.data() + offset, 128);
            uint32_t ref = crc32_reference(buf.data() + offset, 128);
            if (got != ref) {
                std::cerr << "  CRC32 MISMATCH at offset " << offset
                          << ": got 0x" << std::hex << got
                          << " expected 0x" << ref << std::dec << std::endl;
                ok = false;
            }
        }
        if (!ok) { std::cerr << "  FATAL: CRC32 alignment test failed" << std::endl; std::abort(); }
        std::cout << "  Unaligned data test passed (16 offsets)." << std::endl;
    }

    // Test 6: large block (triggers multi-iteration 64-byte folding)
    {
        std::vector<uint8_t> buf(65536);
        for (size_t i = 0; i < buf.size(); ++i)
            buf[i] = static_cast<uint8_t>((i * 37 + i / 256) & 0xFF);

        uint32_t got = file_block_store::crc32(buf.data(), buf.size());
        uint32_t ref = crc32_reference(buf.data(), buf.size());
        if (got != ref) {
            std::cerr << "  FATAL: CRC32(64KB) = 0x" << std::hex << got
                      << " expected 0x" << ref << std::dec << std::endl;
            std::abort();
        }
        std::cout << "  Large block (64KB) test passed." << std::endl;
    }

    std::cout << "  All CRC32 tests passed." << std::endl;
}

static void test_file_block_store() {
    std::cout << "=== File Block Store (write-hole protection) ===" << std::endl;

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

    // Test 1: basic store/fetch with CRC32
    {
        file_block_store store(dir, 0); // threshold=0: flush every write
        std::vector<uint8_t> data = {1, 2, 3, 4, 5, 6, 7, 8};
        assert(store.store(100, 0, data.data(), data.size()));

        std::vector<uint8_t> out;
        assert(store.fetch(100, 0, out));
        assert(out == data);
        assert(store.total_blocks() == 1);
        assert(store.total_groups() == 1);
    }

    // Test 2: reopen and verify data survives (CRC32 validated on scan)
    {
        file_block_store store(dir, 0);
        std::vector<uint8_t> out;
        assert(store.fetch(100, 0, out));
        std::vector<uint8_t> expected = {1, 2, 3, 4, 5, 6, 7, 8};
        assert(out == expected);
        assert(store.total_blocks() == 1);
    }

    // Test 3: overwrite + vacuum
    {
        file_block_store store(dir, 0);
        std::vector<uint8_t> data2 = {10, 20, 30};
        assert(store.store(100, 0, data2.data(), data2.size()));

        std::vector<uint8_t> out;
        assert(store.fetch(100, 0, out));
        assert(out == data2);

        assert(store.vacuum());
        out.clear();
        assert(store.fetch(100, 0, out));
        assert(out == data2);
    }

    // Test 4: multiple groups and blocks
    {
        file_block_store store(dir, 0);
        for (uint64_t g = 200; g < 205; ++g) {
            for (uint32_t b = 0; b < 4; ++b) {
                std::vector<uint8_t> d(64, static_cast<uint8_t>(g + b));
                store.store(g, b, d.data(), d.size());
            }
        }

        assert(store.total_groups() >= 5); // 200-204 + maybe 100 from above

        for (uint64_t g = 200; g < 205; ++g) {
            auto blocks = store.list_blocks(g);
            assert(blocks.size() == 4);
            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 + b));
            }
        }
    }

    // Test 5: remove group
    {
        file_block_store store(dir, 0);
        assert(store.remove_group(200));
        auto blocks = store.list_blocks(200);
        assert(blocks.empty());
    }

    // Test 6: WAL recovery — simulate crash by corrupting end of data file
    {
        // Write some data with buffering (threshold > 0)
        std::filesystem::remove_all(dir);
        {
            file_block_store store(dir, 0);
            std::vector<uint8_t> good = {0xAA, 0xBB, 0xCC, 0xDD};
            store.store(300, 0, good.data(), good.size());
        }

        // Append garbage to simulate a partial/corrupt write
        {
            std::string data_path = dir + "/blocks.bin";
            FILE* f = fopen(data_path.c_str(), "ab");
            assert(f);
            uint8_t garbage[10] = {0xDE, 0xAD, 0xBE, 0xEF, 0x00,
                                    0x00, 0x01, 0x02, 0x03, 0x04};
            fwrite(garbage, 1, sizeof(garbage), f);
            fclose(f);
        }

        // Reopen — should detect corrupt trailing record via CRC and truncate
        {
            file_block_store store(dir, 0);
            std::vector<uint8_t> out;
            assert(store.fetch(300, 0, out));
            std::vector<uint8_t> expected = {0xAA, 0xBB, 0xCC, 0xDD};
            assert(out == expected);
            assert(store.total_blocks() == 1);
        }
    }

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

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

    test_galois();
    test_matrix();
    test_parity_basic();
    test_parity_raid6();
    test_parity_flexible();
    test_split_join();
    test_end_to_end();
    test_auth_store();
    test_auth_hmac();
    test_auth_protocol();
    test_crc32();
    test_file_block_store();

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