#include <balancer/kernel/http2/server/common/http2_headers_fsm.h>

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

#include <util/generic/algorithm.h>
#include <util/generic/deque.h>
#include <util/generic/xrange.h>
#include <util/string/ascii.h>

using namespace std::string_view_literals;

Y_UNIT_TEST_SUITE(HTTP2HeadersFsmTest) {
    using namespace NSrvKernel;
    using namespace NSrvKernel::NHTTP2;

    static TChunkList BuildChunkList(const TVector<TString>& data) {
        TChunkList res;
        for (const auto& item : data) {
            res.Push(item);
        }
        return res;
    }

    template <class Validator>
    void DoTestValidationSuccess(char ch, Validator&& validator) {
        UNIT_ASSERT_C(validator(BuildChunkList({TString(ch)})), TString(ch).Quote());
        UNIT_ASSERT_C(validator(BuildChunkList({"", TString(ch)})), TString(ch).Quote());
        UNIT_ASSERT_C(validator(BuildChunkList({TString(ch), ""})), TString(ch).Quote());
        UNIT_ASSERT_C(validator(BuildChunkList({TString(ch), TString(ch)})), TString(ch).Quote());
    }

    template <class Validator>
    void DoTestItemValidationFailure(TVector<TString> data, Validator&& validator) {
        UNIT_ASSERT(!validator(BuildChunkList(data)));
        {
            auto data1 = data;
            data1.push_back("");
            UNIT_ASSERT(!validator(BuildChunkList(data1)));
        }
        {
            auto data2 = data;
            data2.insert(data2.begin(), "");
            UNIT_ASSERT(!validator(BuildChunkList(data2)));
        }
    }

    template <class Validator>
    void DoTestValidationFailure(char ch, Validator&& validator, bool onlyEnd = false) {
        DoTestItemValidationFailure({TString(ch)}, validator);

        DoTestItemValidationFailure({TString::Join(ch, "a")}, validator);
        DoTestItemValidationFailure({TString(ch), "a"}, validator);

        DoTestItemValidationFailure({TString::Join("a", ch)}, validator);
        DoTestItemValidationFailure({"a", TString(ch)}, validator);

        DoTestItemValidationFailure({TString::Join(ch, "aa")}, validator);
        DoTestItemValidationFailure({TString::Join(ch, "a"), "a"}, validator);
        DoTestItemValidationFailure({TString(ch), "aa"}, validator);

        DoTestItemValidationFailure({TString::Join("aa", ch)}, validator);
        DoTestItemValidationFailure({"a", TString::Join("a", ch)}, validator);
        DoTestItemValidationFailure({"aa", TString(ch)}, validator);


        if (!onlyEnd) {
            DoTestItemValidationFailure({TString::Join("a", ch, "a")}, validator);
            DoTestItemValidationFailure({TString::Join("a", ch), "a"}, validator);
            DoTestItemValidationFailure({"a", TString::Join(ch, "a")}, validator);
            DoTestItemValidationFailure({"a", TString(ch), "a"}, validator);

            DoTestItemValidationFailure({TString::Join("a", ch, "a"), "a"}, validator);
            DoTestItemValidationFailure({"a", TString::Join("a", ch, "a")}, validator);
        }
    }

    Y_UNIT_TEST(TestRequestNameValidation) {
        const TStringBuf validChars = "!#$%&'*+._`|~^-";

        UNIT_ASSERT(!IsRequestHeaderNameValid(BuildChunkList({})));
        UNIT_ASSERT(!IsRequestHeaderNameValid(BuildChunkList({"", "", ""})));

        for (int ch : xrange(0, 256)) {
            if (IsAsciiAlnum((char)ch) && !IsAsciiUpper((char)ch) || validChars.find((char)ch) != TStringBuf::npos) {
                DoTestValidationSuccess(ch, IsRequestHeaderNameValid);
            } else {
                DoTestValidationFailure(ch, IsRequestHeaderNameValid);
            }
        }
    }

    Y_UNIT_TEST(TestRequestValueValidation) {
        constexpr TStringBuf invalidChars = "\0\r\n"sv;
        constexpr TStringBuf invalidCharsEnd = "\0\r\n\t "sv;

        for (int ch : xrange(0, 256)) {
            if (invalidChars.find((char)ch) != TStringBuf::npos) {
                DoTestValidationFailure(ch, IsRequestHeaderValueValid, false);
            } else if (invalidCharsEnd.find(ch) != TStringBuf::npos) {
                DoTestValidationFailure(ch, IsRequestHeaderValueValid, true);
            } else {
                DoTestValidationSuccess(ch, IsRequestHeaderValueValid);
            }
        }
    }
}
