#include <balancer/production/x/collect_cookie_stats/lib/cookie_extractor.h>

#include <library/cpp/yson/node/node.h>
#include <library/cpp/yson/node/node_io.h>

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

Y_UNIT_TEST_SUITE(CookieExtractor) {
    struct TExpectedCookie {
        TString Name;
        ui32 ValueLenMax = 0;
        TString ValueSet;
        ui32 NoValueCount = 0;
        ui32 ReqCount = 1;
        ui32 TotalCount = 1;
        ui32 InReqCountMax = 1;

    public:
        void Check(const NCookie::TCookie& cookie, ui32 called) const {
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetName(), Name,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetValueLenMax(), ValueLenMax,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetValueSet(), ValueSet,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetNoValueCount(), NoValueCount,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetReqCount(), ReqCount,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetTotalCount(), TotalCount,
                called << ": " << cookie.ShortDebugString());
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetInReqCountMax(), InReqCountMax,
                called << ": " << cookie.ShortDebugString());
        }
    };

    Y_UNIT_TEST(TestExtractCookie) {
        using namespace NCookie;
        const TVector<TExpectedCookie> expected {
            {.Name="x", .NoValueCount=1,},
            {"yandexuid", 21, "!a\\xFF",},
            {"yandex_gid", 3, "a",},
            {.Name="x", .NoValueCount=1,},
            {"yandexuid", 19, "a",},
            {"mda", 1, "a"},
            {"yandex_gid", 4, "@a",},
            {"L", 105, "+.a"},
            {"yabs-frequency",  41, "/a"},
            {"my",  8, "a"},
            {"_ym_isad",  1, "a"},
            {"zm",  77, "%-._a"},
            {"yc",  32, "%.a"},
            {"i",  92, "+/=a"},
            {"Session_id",  98, ".:a|"},
            {"sessionid2",  98, ".:_a|"},
            {"yandex_login",  12, "a"},
            {"ys",  67, "#-._a"},
            {"yp",  279, "#%.:a"},
            {"bltsr",  1, "a"},
            {"skid",  19, "a"},
            {"_ym_visorc_160656",  1, "a"},
            {"_ym_visorc_45411513",  1, "a"},
            {"_ym_uid",  18, "a"},
            {"_ym_d",  10, "a"},
        };

        const auto row = NYT::NodeFromYsonString(R"=({
    "workflow" = " [report u:https [regexp_path default [log_headers <::Host:wg6-bro.yandex.ru::> <::CookieMeta:x;yandexuid=21:!a\\xFF;yandex_gid=3:a;::> <::CookieMeta:x;yandexuid=19:a;mda=1:a;yandex_gid=4:@a;L=105:+.a;yabs-frequency=41:/a;my=8:a;_ym_isad=1:a;zm=77:%-._a;yc=32:%.a;i=92:+/=a;Session_id=98:.:a|;sessionid2=98:.:_a|;yandex_login=12:a;ys=67:#-._a;yp=279:#%.:a;bltsr=1:a;skid=19:a;_ym_visorc_160656=1:a;_ym_visorc_45411513=1:a;_ym_uid=18:a;_ym_d=10:a::> [report u:any_requests_to_sas [proxy sas1-3914-sas-portal-any-stable-27240.gencfg-c.yandex.net:27240 0.043693s/0.044027s 0/13060 succ 404]]]]]";
})=");

        ui32 called = 0;
        TCookieExtractor extractor;
        extractor.Apply(row, [&](const TCookie& cookie) {
            UNIT_ASSERT_LT(called, expected.size());
            expected[called].Check(cookie, called);
            called += 1;
        });

        UNIT_ASSERT_VALUES_EQUAL(called, expected.size());
    }

    struct TExpectedSetCookie {
        TString Name;
        ui32 ValueLen = 0;
        TString ValueSet;
        bool NoValue = false;

        ui32 InReqCount = 0;
        TString ReqPath;
        TString ReqHost;

        TString Domain;
        TString Path;
        TString SameSite;
        TString Secure;
        TString HttpOnly;
        TString MaxAge;
        TString Expires;
        TString OtherAttrs;

        void Check(const NCookie::TSetCookie& cookie, ui32 called) const {
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetName(), Name, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetValueLen(), ValueLen, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetValueSet(), ValueSet, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetNoValue(), NoValue, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetInReqCount(), InReqCount, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetReqPath(), ReqPath, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetReqHost(), ReqHost, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetDomain(), Domain, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetPath(), Path, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetSameSite(), SameSite, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetSecure(), Secure, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetHttpOnly(), HttpOnly, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetMaxAge(), MaxAge, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetExpires(), Expires, called);
            UNIT_ASSERT_VALUES_EQUAL_C(cookie.GetOtherAttrs(), OtherAttrs, called);
        }
    };

    Y_UNIT_TEST(TestExtractSetCookie) {
        using namespace NCookie;
        const TVector<TExpectedCookie> expectedCookie {
            {"yandexuid", 19, "a",},
            {"yandexuid", 19, "a",},
            {"ys", 57, "-.a"},
            {"i", 92, "/=a"},
        };
        const TVector<TExpectedSetCookie> expectedSetCookie {
            {
                .Name="xxx",
                .NoValue=true,
                .ReqPath="/ick/r",
                .ReqHost="yandex.ru"
            }, {
                .Name="xxx",
                .ValueLen=1,
                .ValueSet="a",
                .ReqPath="/ick/r",
                .ReqHost="yandex.ru"
            }, {
                .Name="i",
                .ValueLen=93,
                .ValueSet="!=a",
                .InReqCount=1,
                .ReqPath="/ick/r",
                .ReqHost="yandex.ru",
                .Domain="domain=.yandex.ru",
                .Path="path=/",
                .SameSite="samesite=None",
                .Secure="secure",
                .HttpOnly="httponly",
                .MaxAge="max-age=100500",
                .Expires="expires=Fri, 18-Jan-2030 16:38:25 GMT",
                .OtherAttrs="blabla"
            },
        };

        const auto row = NYT::NodeFromYsonString(R"=({
"host" = "\"yandex.ru\"";
"query" = "\"POST /ick/r HTTP/2\"";
"workflow" = " [report u:service_total u:service_total_ru u:service_total_h2 [regexp yandex [h100 [regexp default [log_headers [log_headers <::X-Yandex-TCP-Info:v=2; rtt=0.058334s; rttvar=0.011083s; snd_cwnd=84; total_retrans=0::> <::Y-Balancer-Experiments:no::> <::X-Yandex-IP:77.88.55.60::> <::X-Req-Id:1579624705180679-10917974443985083843::> <::user-agent:Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/79.0.3945.117 Safari/537.36::> <::host:yandex.ru::> <::Cookie:yandexuid=9719858541579624704::> <::CookieMeta:yandexuid=19:a;yandexuid=19:a;ys=57:-.a;i=92:/=a::> [regexp_path default [report u:default [h100 [cutter [antirobot [sub_antirobot [report u:antirobot [proxy sas2-9244.search.yandex.net:13512 0.000526s/0.000558s/connect=0.000000s 1525/187 succ 200]] not_robot] [sub_search [report u:default_requests_to_knoss_localdc [proxy sas1-3521-a26-sas-l7-balancer-k-2f7-7734.gencfg-c.yandex.net:7734 0.012268s/0.012433s/connect=0.003343s 30/391 succ 200]]]]]]]]] <::SetCookieMeta:xxx::> <::SetCookieMeta:xxx=1:a::> <::SetCookieMeta:i=93:!=a;blabla;Expires=Fri, 18-Jan-2030 16:38:25 GMT;Domain=.yandex.ru;Path=/;Secure;HttpOnly;SameSite=None;Max-Age=100500::>]]]]]";
})=");

        ui32 calledCookie = 0;
        ui32 calledSetCookie = 0;
        TSetCookieExtractor extractor;
        auto hostPath = extractor.Apply(row, [&](const TCookie& cookie) {
            UNIT_ASSERT_LT(calledCookie, expectedCookie.size());
            expectedCookie[calledCookie].Check(cookie, calledCookie);
            calledCookie += 1;
        }, [&](const TSetCookie& setCookie) {
            UNIT_ASSERT_LT(calledSetCookie, expectedSetCookie.size());
            expectedSetCookie[calledSetCookie].Check(setCookie, calledSetCookie);
            calledSetCookie += 1;
        });
        UNIT_ASSERT_VALUES_EQUAL(calledCookie, expectedCookie.size());
        UNIT_ASSERT_VALUES_EQUAL(calledSetCookie, expectedSetCookie.size());
        UNIT_ASSERT_VALUES_EQUAL(hostPath.Host, "yandex.ru");
        UNIT_ASSERT_VALUES_EQUAL(hostPath.Path, "/ick/r");
    }
}
