#include <balancer/kernel/memory/strip.h>

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

#include <util/string/join.h>
#include <util/generic/algorithm.h>

Y_UNIT_TEST_SUITE(ChunkStripTest) {
    void DoTestChunkStrip(std::function<bool(NSrvKernel::TChunk&, char)> str, TStringBuf before, TStringBuf after) {
        using namespace NSrvKernel;

        TChunkPtr ch = NewChunkNonOwning(before);
        const bool res = str(*ch, ' ');

        UNIT_ASSERT_VALUES_EQUAL_C(ch->AsStringBuf(), after, before);
        UNIT_ASSERT_VALUES_EQUAL_C(before != after, res, before);
    }

    void DoTestListStrip(std::function<bool(NSrvKernel::TChunkList&, char)> str, TVector<TStringBuf> before, TVector<TStringBuf> after) {
        using namespace NSrvKernel;

        TChunkList lst;
        for (auto b : before) {
            lst.PushNonOwning(b);
        }

        before.erase(std::remove_if(before.begin(), before.end(), std::mem_fn(&TStringBuf::empty)), before.end());
        after.erase(std::remove_if(after.begin(), after.end(), std::mem_fn(&TStringBuf::empty)), after.end());

        const bool res = str(lst, ' ');

        TVector<TStringBuf> result;
        std::transform(lst.ChunksBegin(), lst.ChunksEnd(), std::back_inserter(result),
                       std::mem_fn(&TChunk::AsStringBuf));

        UNIT_ASSERT_VALUES_EQUAL_C(result, after, JoinSeq(";", before).Quote());
        UNIT_ASSERT_VALUES_EQUAL_C(before != after, res, JoinSeq(";", before).Quote());
    }

    Y_UNIT_TEST(TestChunkLStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunk&, char);
        TStrp strp = &LStrip;
        DoTestChunkStrip(strp, TStringBuf(), TStringBuf());
        DoTestChunkStrip(strp, "", "");
        DoTestChunkStrip(strp, " ", "");
        DoTestChunkStrip(strp, "  ", "");
        DoTestChunkStrip(strp, " a ", "a ");
        DoTestChunkStrip(strp, "aa ", "aa ");
    }

    Y_UNIT_TEST(TestChunkRStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunk&, char);
        TStrp strp = &RStrip;
        DoTestChunkStrip(strp, TStringBuf(), TStringBuf());
        DoTestChunkStrip(strp, "", "");
        DoTestChunkStrip(strp, " ", "");
        DoTestChunkStrip(strp, "  ", "");
        DoTestChunkStrip(strp, " a ", " a");
        DoTestChunkStrip(strp, " aa", " aa");
    }

    Y_UNIT_TEST(TestChunkStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunk&, char);
        TStrp strp = &Strip;
        DoTestChunkStrip(strp, TStringBuf(), TStringBuf());
        DoTestChunkStrip(strp, "", "");
        DoTestChunkStrip(strp, " ", "");
        DoTestChunkStrip(strp, "  ", "");
        DoTestChunkStrip(strp, " a ", "a");
        DoTestChunkStrip(strp, " aa", "aa");
    }

    Y_UNIT_TEST(TestListLStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunkList&, char);
        TStrp strp = &LStrip;
        DoTestListStrip(strp, {}, {});
        DoTestListStrip(strp, {TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {" ", TStringBuf(), " "}, {});
        DoTestListStrip(strp, {TStringBuf(), " ", TStringBuf()}, {});
        DoTestListStrip(strp, {"a"}, {"a"});
        DoTestListStrip(strp, {" a"}, {"a"});
        DoTestListStrip(strp, {" a "}, {"a "});
        DoTestListStrip(strp, {"a "}, {"a "});
        DoTestListStrip(strp, {" ", " a ", " "}, {"a ", " "});
        DoTestListStrip(strp, {" ", "a ", " "}, {"a ", " "});
    }

    Y_UNIT_TEST(TestListRStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunkList&, char);
        TStrp strp = &RStrip;
        DoTestListStrip(strp, {}, {});
        DoTestListStrip(strp, {TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {" ", TStringBuf(), " "}, {});
        DoTestListStrip(strp, {TStringBuf(), " ", TStringBuf()}, {});
        DoTestListStrip(strp, {"a"}, {"a"});
        DoTestListStrip(strp, {"a "}, {"a"});
        DoTestListStrip(strp, {" a "}, {" a"});
        DoTestListStrip(strp, {" a"}, {" a"});
        DoTestListStrip(strp, {" ", " a ", " "}, {" ", " a"});
        DoTestListStrip(strp, {" ", " a", " "}, {" ", " a"});
    }

    Y_UNIT_TEST(TestListStrip) {
        using namespace NSrvKernel;
        using TStrp = bool(*)(TChunkList&, char);
        TStrp strp = &Strip;
        DoTestListStrip(strp, {}, {});
        DoTestListStrip(strp, {TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {TStringBuf(), TStringBuf(), TStringBuf()}, {});
        DoTestListStrip(strp, {" ", TStringBuf(), " "}, {});
        DoTestListStrip(strp, {TStringBuf(), " ", TStringBuf()}, {});
        DoTestListStrip(strp, {"a"}, {"a"});
        DoTestListStrip(strp, {" a "}, {"a"});
        DoTestListStrip(strp, {" a"}, {"a"});
        DoTestListStrip(strp, {"a "}, {"a"});
        DoTestListStrip(strp, {" ", " a ", " "}, {"a"});
        DoTestListStrip(strp, {" ", " a", " "}, {"a"});
        DoTestListStrip(strp, {" ", "a ", " "}, {"a"});
    }
}
