#include "erasure.h"

#include <library/cpp/testing/unittest/registar.h>

#include <library/cpp/digest/crc32c/crc32c.h>

extern "C" {
    #include <contrib/libs/jerasure/jerasure.h>
    #include <contrib/libs/jerasure/reed_sol.h>
}

Y_UNIT_TEST_SUITE(Erasure) {
    Y_UNIT_TEST(JErasureNotBreaksMainParts) {
        ui32 mainNumber = 3, parityNumber = 2, blockSize = 100;
        int *matrix = reed_sol_vandermonde_coding_matrix(mainNumber, parityNumber, 8);
        TVector<char *> parts_ptrs(mainNumber + parityNumber);
        TString parts[mainNumber + parityNumber];
        ui32 sums[mainNumber];
        for (ui32 i = 0; i < mainNumber + parityNumber; ++i) {
            parts[i].resize(blockSize);
            parts_ptrs[i] = (char *)parts[i].data();
            if (i < mainNumber) {
                for (size_t j = 0; j < blockSize; ++j) {
                    parts[i][j] = (char)(rand() % 256);
                }
                sums[i] = Crc32c(parts[i].begin(), parts[i].size());
            }
        }

        jerasure_matrix_encode(mainNumber, parityNumber, 8, matrix, parts_ptrs.begin(), parts_ptrs.begin() + mainNumber, blockSize); // work, pls

        for (ui32 i = 0; i < mainNumber; ++i) {
            UNIT_ASSERT(sums[i] == Crc32c(parts[i].begin(), parts[i].size()));
        }
    }

    Y_UNIT_TEST(JustWorks) {
        NSaas::TErasureHelper erasure;
        TString testString("Just simple string which should fail this lib.");

        auto splitted = erasure.SplitToParts(testString);

        auto recoveredString = erasure.Merge(splitted);

        UNIT_ASSERT_EQUAL(testString, recoveredString);
    }

    Y_UNIT_TEST(JustWorksWithCorrupting) {
        NSaas::TErasureHelper erasure;
        TString testString("Just simple string which should fail this lib.");

        auto splitted = erasure.SplitToParts(testString);
        // Corrupting...
        *splitted[0].MutableData() = "";

        auto recoveredString = erasure.Merge(splitted);

        UNIT_ASSERT_EQUAL(testString, recoveredString);
    }

    Y_UNIT_TEST(IncreasedSizeStrings) {
        NSaas::TErasureHelper erasure;
        TString testString;
        ui32 len = 10000;
        testString.resize(len);

        for (ui32 i = 0; i < testString.size(); ++i) {
            testString[i] = 'a' + i % (int)('z' - 'a');
        }

        auto splitted = erasure.SplitToParts(testString);
        // Corrupting...
        *splitted[2].MutableData() = "";
        *splitted[4].MutableData() = "";

        auto recoveredString = erasure.Merge(splitted);

        UNIT_ASSERT_EQUAL(testString, recoveredString);
    }

    Y_UNIT_TEST(TooManyErasures) {
        NSaas::TErasureHelper erasure;
        TString testString;
        ui32 len = 10000;
        testString.resize(len);

        for (ui32 i = 0; i < testString.size(); ++i) {
            testString[i] = 'a' + i % (int)('z' - 'a');
        }

        auto splitted = erasure.SplitToParts(testString);
        // Corrupting...
        *splitted[2].MutableData() = "";
        *splitted[4].MutableData() = "";
        *splitted[1].MutableData() = "";
        bool catched = false;
        try {
            auto recoveredString = erasure.Merge(splitted);
        } catch (...) {
            // everything is alright
            catched = true;
        }
        UNIT_ASSERT(catched);
    }

    Y_UNIT_TEST(NotAllPartsReceived) {
        NSaas::TErasureHelper erasure;
        TString testString;
        ui32 len = 10000;
        testString.resize(len);

        for (ui32 i = 0; i < testString.size(); ++i) {
            testString[i] = 'a' + i % (int)('z' - 'a');
        }

        auto splitted = erasure.SplitToParts(testString);
        // Corrupting...
        std::reverse(splitted.begin(), splitted.end());
        splitted.pop_back();

        auto recoveredString = erasure.Merge(splitted);

        UNIT_ASSERT_EQUAL(testString, recoveredString);
    }
}