#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 <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 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_file_block_store();

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