#include <library/cpp/testing/unittest/registar.h>
#include <balancer/kernel/custom_io/double.h>
#include <balancer/kernel/custom_io/ut/test_common.h>
#include <util/random/random.h>

Y_UNIT_TEST_SUITE(TCustomIoDoubleTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA = "supermegatesthellocarzil";
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);

    Y_UNIT_TEST(TestDoubleInput) {
        size_t position = 0;
        auto mock = MakeInputMock([&position](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (position < DATA.Size()) {
                lst.PushNonOwning(DATA.SubStr(position++, 1));
            }
            return TError{};
        });

        TDoubleInputMaster master{mock, 30LL << 20};
        TDoubleInputSlave slave0{master, 0};
        TDoubleInputSlave slave1{master, 1};

        TChunkList result0;
        TChunkList result1;

        while (position < DATA.Size()) {
            TChunkList lst;
            if (RandomNumber<bool>()) {
                auto err = slave0.Recv(lst, DLINE);
                UNIT_ASSERT(!err);
                result0.Append(std::move(lst));
            } else {
                auto err = slave1.Recv(lst, DLINE);
                UNIT_ASSERT(!err);
                result1.Append(std::move(lst));
            }
        }

        TChunkList lst0;
        auto err0 = slave0.Recv(lst0, DLINE);
        UNIT_ASSERT(!err0);
        result0.Append(std::move(lst0));

        TChunkList lst1;
        auto err1 = slave1.Recv(lst1, DLINE);
        UNIT_ASSERT(!err1);
        result1.Append(std::move(lst1));

        UNIT_ASSERT(Union(result0)->AsStringBuf() == DATA);
        UNIT_ASSERT(Union(result1)->AsStringBuf() == DATA);
    }

    Y_UNIT_TEST(TestDoubleInputLimit) {
        const size_t LIMIT = DATA.Size() / 2;

        size_t position = 0;
        auto mock = MakeInputMock([&position](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (position < DATA.Size()) {
                lst.PushNonOwning(DATA.SubStr(position++, 1));
            }
            return TError{};
        });

        TDoubleInputMaster master{mock, LIMIT};
        TDoubleInputSlave slave0{master, 0};
        TDoubleInputSlave slave1{master, 1};

        TChunkList result0;
        TChunkList result1;

        while (position < DATA.Size()) {
            TChunkList lst;
            auto err = slave0.Recv(lst, DLINE);
            UNIT_ASSERT(!err);
            result0.Append(std::move(lst));
        }

        TChunkList lst0;
        auto err0 = slave0.Recv(lst0, DLINE);
        UNIT_ASSERT(!err0);
        result0.Append(std::move(lst0));

        TChunkList lst1;
        auto err1 = slave1.Recv(lst1, DLINE);
        UNIT_ASSERT(!err1);
        result1.Append(std::move(lst1));

        UNIT_ASSERT(Union(result0)->AsStringBuf() == DATA);
        UNIT_ASSERT(Union(result1)->AsStringBuf().size() <= LIMIT + 1);
    }

    Y_UNIT_TEST(TestBufferNoOverflow) {
        const TString MB_DATA(1'000'000, 'A');
        const size_t ITERATIONS = 2;

        size_t currentIteration = 0;

        auto mock = MakeInputMock([&MB_DATA, &currentIteration](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (currentIteration < ITERATIONS) {
                lst.PushNonOwning(TStringBuf(MB_DATA));
                ++currentIteration;
            }
            return TError{};
        });

        TDoubleInputMaster master{mock, 30LL << 20};
        TDoubleInputSlave slave0{master, 0};
        TDoubleInputSlave slave1{master, 1};

        size_t total0 = 0;
        size_t total1 = 0;

        while (currentIteration < ITERATIONS) {
            TChunkList lst0;
            auto err0 = slave0.Recv(lst0, DLINE);
            UNIT_ASSERT(!err0);
            total0 += lst0.size();

            TChunkList lst1;
            auto err1 = slave1.Recv(lst1, DLINE);
            UNIT_ASSERT(!err1);
            total1 += lst1.size();
        }

        TChunkList lst0;
        auto err0 = slave0.Recv(lst0, DLINE);
        UNIT_ASSERT(!err0);
        total0 += lst0.size();

        TChunkList lst1;
        auto err1 = slave1.Recv(lst1, DLINE);
        UNIT_ASSERT(!err1);
        total1 += lst1.size();

        UNIT_ASSERT_VALUES_EQUAL(total0, ITERATIONS * MB_DATA.size());
        UNIT_ASSERT_VALUES_EQUAL(total1, ITERATIONS * MB_DATA.size());
    }
}
