#include <balancer/kernel/custom_io/ut/test_common.h>
#include <balancer/kernel/custom_io/iterator.h>
#include <balancer/kernel/helpers/errors.h>
#include <library/cpp/testing/unittest/registar.h>

Y_UNIT_TEST_SUITE(TCustomIoIteratorTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA = "test";
    const TStringBuf OTHER_DATA = "data";
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);

    void DoRead(TStreamIterator& it) {
        char buf[DATA.size() + 1];
        buf[DATA.size()] = '\0';
        auto sz = it.ReadSome(buf, DATA.size(), DLINE);
        UNIT_ASSERT(!sz);
        size_t res;
        auto error = sz.AssignTo(res);
        UNIT_ASSERT(!error);
        UNIT_ASSERT_EQUAL(res, DATA.size());
        UNIT_ASSERT_EQUAL(TString(buf), TString(DATA));

        for (size_t i = 0; i < DATA.size(); ++i) {
            char c;
            auto next = it.Next(c, DLINE);
            UNIT_ASSERT(!next);
            bool res;
            error = next.AssignTo(res);
            UNIT_ASSERT(!error);
            UNIT_ASSERT(res);
            UNIT_ASSERT_EQUAL(c, DATA[i]);
        }

        TChunkList chunk;
        auto rd = it.Read(chunk, DATA.size(), DLINE);
        UNIT_ASSERT(!rd);
        UNIT_ASSERT_EQUAL(chunk, DATA);

        auto err = it.Read(buf, DATA.size(), DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT_EQUAL(TString(buf), TString(DATA));

    }

    Y_UNIT_TEST(TestStreamIterator) {
        auto mock = MakeInputMock([](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            lst.PushNonOwning(DATA);
            return TError{};
        });

        TStreamIterator it1(&mock);

        TChunkList data(NewChunkNonOwning(OTHER_DATA));
        it1.UnGet(data);
        UNIT_ASSERT(data.Empty());

        char buf[OTHER_DATA.size() + 1];
        buf[OTHER_DATA.size()] = '\0';
        auto sz = it1.ReadSome(buf, OTHER_DATA.size(), DLINE);
        UNIT_ASSERT(!sz);
        size_t res;
        auto error = sz.AssignTo(res);
        UNIT_ASSERT(!error);
        UNIT_ASSERT_EQUAL(res, OTHER_DATA.size());
        UNIT_ASSERT_EQUAL(TString(buf), TString(OTHER_DATA));
        UNIT_ASSERT(!it1.HasBufferedData());

        auto result = it1.Peek(DLINE);
        UNIT_ASSERT(!result);
        bool peeked;
        error = result.AssignTo(peeked);
        UNIT_ASSERT(peeked);
        UNIT_ASSERT(it1.HasBufferedData());
        DoRead(it1);

        TStreamIterator it2(std::move(it1));
        DoRead(it2);

        TChunkList lst(NewChunkNonOwning(DATA));
        TStreamIterator it3(std::move(lst));
        UNIT_ASSERT(lst.Empty());
        auto err1 = it3.Read(buf, DATA.size(), DLINE);
        UNIT_ASSERT(!err1);
        UNIT_ASSERT_EQUAL(TString(buf), TString(DATA));
        auto err2 = it3.ReadSome(buf, DATA.size(), DLINE);
        UNIT_ASSERT(!err2);
        error = err2.AssignTo(res);
        UNIT_ASSERT(!res);

    }

    void DoReadWithError(TStreamIterator& it) {
        char buf[DATA.size() + 1];
        auto err = it.Read(buf, DATA.size(), DLINE);
        UNIT_ASSERT(err.GetAs<TSystemError>());
    }

    Y_UNIT_TEST(TestStreamIteratorWithError) {
        auto mock = MakeInputMock([](TChunkList&, TInstant deadline) {
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return Y_MAKE_ERROR(TSystemError{});
        });

        TStreamIterator it1(&mock);
        DoReadWithError(it1);

        TStreamIterator it2(std::move(it1));
        DoReadWithError(it2);
    }

    Y_UNIT_TEST(TestStreamIteratorWithTStreamError) {
        TChunkList lst;
        TStreamIterator it(std::move(lst));
        char buf[DATA.size() + 1];
        auto error = it.Read(buf, DATA.size(), DLINE);
        UNIT_ASSERT(error);
        UNIT_ASSERT_UNEQUAL(error.GetAs<TStreamError>(), nullptr);
        UNIT_ASSERT_EQUAL(error.GetAs<TStreamError>()->GetUnreadLength(), DATA.size());
    }
};
