#include <balancer/modules/debug/debug_io.h>

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

#include <util/generic/xrange.h>

Y_UNIT_TEST_SUITE(TTestDebugIo) {
    using namespace NModDebug;


    static auto GenChunkList(const TVector<TString>& chunks) {
        NSrvKernel::TChunkList res;
        for (auto&& ch : chunks) {
            res.Push(ch);
        }
        return res;
    }


    struct TNextItem : public TIntrusiveListItem<TNextItem> {
        TInstant Time = TInstant::Now();
        TChunkList Data;

        explicit TNextItem(TChunkList lst)
            : Data(std::move(lst))
        {}
    };


    struct TTestArgs : public TDebugIoParams {
        TChunkList Data;

        TTestArgs(TVector<TString> data, TDuration dur, size_t pSize, size_t health)
            : TDebugIoParams(dur, pSize, health)
            , Data(GenChunkList(data))
        {}
    };


    struct TTestItems final : IIoInput, IHttpOutput {
        TTestArgs Args;
        TInstant Start = TInstant::Now();
        TIntrusiveListWithAutoDelete<TNextItem, TDelete> Result;
        TError Err;
        bool Eof = false;

    public:
        TTestItems(TTestArgs args)
            : Args(std::move(args))
        {}

        TString Str() const {
            TStringBuilder str;
            for (const auto& item : Result) {
                str << item.Data;
            }
            return str;
        }

        TError DoRecv(NSrvKernel::TChunkList& lst, TInstant) noexcept override {
            lst.Append(std::move(Args.Data));
            return {};
        }

        void DoRunRecv(TCont* cont) noexcept {
            TDebugIn in{*this, *cont->Executor(), Args};

            while (!Eof) {
                TChunkList res;
                Err = in.Recv(res, TInstant::Max());
                if (Err) {
                    break;
                }
                if (!res.Empty()) {
                    Result.PushBack(new TNextItem{std::move(res)});
                } else {
                    Eof = true;
                }
            }
        }

        TError DoSendHead(TResponse&& response, const bool forceClose, TInstant) override {
            Y_UNUSED(response, forceClose);
            // TODO(tender-bum): SEND_HEAD
            return {};
        }

        TError DoSend(NSrvKernel::TChunkList lst, TInstant) noexcept override {
            if (!lst.Empty()) {
                Result.PushBack(new TNextItem{std::move(lst)});
            } else {
                Eof = true;
            }
            return {};
        }

        TError DoSendTrailers(THeaders&&, TInstant) override {
            return {};
        }

        void DoRunSend(TCont* cont) noexcept {
            TDebugOut out{*this, *cont->Executor(), Args};

            Err = out.Send(std::move(Args.Data), TInstant::Max());
            if (!Err) {
                Err = out.SendEof(TInstant::Max());
            }
        }

        void RunRecv() {
            TContExecutor exec{50000};
            exec.Execute<TTestItems, &TTestItems::DoRunRecv>(this);
        }

        void RunSend() {
            TContExecutor exec{50000};
            exec.Execute<TTestItems, &TTestItems::DoRunSend>(this);
        }

        TString DumpResult() {
            TStringBuilder str;
            for (const auto& item : Result) {
                str << "[" << item.Data << "]@" << (item.Time - Start).MilliSeconds() << ",";
            }
            return str;
        }
    };


    void DoRunDebugIoTest(TTestArgs args, std::function<void(TTestItems&)> run) {
        TString expected = TStringBuilder() << args.Data;
        TTestItems items(std::move(args));
        run(items);

        UNIT_ASSERT_VALUES_EQUAL(items.Args.Data.ChunksCount(), 0);
        UNIT_ASSERT_VALUES_EQUAL(items.Str(), expected.substr(0, args.FailAfter));

        UNIT_ASSERT_VALUES_EQUAL_C(
            items.Result.Size(), (std::min(expected.size(), args.FailAfter) + args.Size - 1) / args.Size, items.DumpResult());

        for (const auto& item : items.Result) {
            if (args.Delay) {
                UNIT_ASSERT((item.Time - items.Start) > args.Delay / 2);
            }
        }
        if (expected.size() < args.FailAfter) {
            UNIT_ASSERT(items.Eof);
        } else {
            UNIT_ASSERT(items.Err);
        }
    }

    void DoTestDebugIn(TTestArgs args) {
        DoRunDebugIoTest(std::move(args), [](TTestItems& items) { items.RunRecv(); });
    }

    void DoTestDebugIn(TDuration dur) {
        DoTestDebugIn({{}, dur, 1, 1});
        DoTestDebugIn({{""}, dur, 1, 1});
        DoTestDebugIn({{"a"}, dur, 1, 2});
        DoTestDebugIn({{"ab"}, dur, 1, 3});
        DoTestDebugIn({{"a", "", "b"}, dur, 1, 3});
        DoTestDebugIn({{"abc"}, dur, 1, 4});
        DoTestDebugIn({{"ab", "c"}, dur, 1, 4});
        DoTestDebugIn({{"a", "b", "c"}, dur, 1, 4});
        DoTestDebugIn({{"ab", "cd"}, dur, 1, 5});
        DoTestDebugIn({{"a", "bc", "d"}, dur, 1, 5});

        for (size_t pSz : xrange(1, 6)) {
            DoTestDebugIn({{"abcd"}, dur, pSz, 5});
            DoTestDebugIn({{"a", "bc", "d"}, dur, pSz, 5});
        }
    }

    Y_UNIT_TEST(TestDebugInNoDelay) {
        DoTestDebugIn(TDuration());
    }

    Y_UNIT_TEST(TestDebugInDelay) {
        DoTestDebugIn(TDuration::MilliSeconds(50));
    }

    void DoTestDebugOut(TTestArgs args) {
        DoRunDebugIoTest(std::move(args), [](TTestItems& items) {
            items.RunSend();
        });
    }

    void DoTestDebugOut(TDuration dur) {
        DoTestDebugOut({{}, dur, 1, 1});
        DoTestDebugOut({{""}, dur, 1, 1});
        DoTestDebugOut({{"a"}, dur, 1, 2});
        DoTestDebugOut({{"ab"}, dur, 1, 3});
        DoTestDebugOut({{"a", "", "b"}, dur, 1, 3});
        DoTestDebugOut({{"abc"}, dur, 1, 4});
        DoTestDebugOut({{"ab", "c"}, dur, 1, 4});
        DoTestDebugOut({{"a", "b", "c"}, dur, 1, 4});
        DoTestDebugOut({{"ab", "cd"}, dur, 1, 5});
        DoTestDebugOut({{"a", "bc", "d"}, dur, 1, 5});

        for (size_t pSz : xrange(1, 6)) {
            DoTestDebugOut({{"abcd"}, dur, pSz, 5});
            DoTestDebugOut({{"a", "bc", "d"}, dur, pSz, 5});
        }

        DoTestDebugOut({{}, dur, 1, 0});
        DoTestDebugOut({{"a"}, dur, 1, 1});
        DoTestDebugOut({{"a"}, dur, 1, 0});
        DoTestDebugOut({{"ab"}, dur, 1, 2});
        DoTestDebugOut({{"ab"}, dur, 1, 1});
        DoTestDebugOut({{"ab"}, dur, 1, 0});
    }

    Y_UNIT_TEST(TestDebugOutNoDelay) {
        DoTestDebugOut(TDuration());
    }

    Y_UNIT_TEST(TestDebugOutDelay) {
        DoTestDebugOut(TDuration::MilliSeconds(50));
    }
}
