#include "sharding.h"
#include <saas/protos/rtyserver.pb.h>
#include <saas/library/sharding/rules/kps.h>
#include <saas/library/sharding/rules/erasure.h>
#include <saas/library/sharding/rules/urlhash.h>

#include <library/cpp/testing/unittest/registar.h>
#include <library/cpp/logger/global/global.h>
#include <util/digest/city.h>
#include <util/random/random.h>
#include <util/generic/algorithm.h>
#include <util/generic/serialized_enum.h>
#include <util/string/strip.h>

using namespace NSaas;

void InitLog() {
    if (!GlobalLogInitialized())
        DoInitGlobalLog("cout", 7, false, false);
}

Y_UNIT_TEST_SUITE(TRtyConverterTests) {
    Y_UNIT_TEST(InitializeFromBaseUrl) {
        TErasureUrlConverter converter;
        UNIT_ASSERT(converter.InitializeFromBaseUrl("example.com"));
        UNIT_ASSERT(converter.GetBaseUrl() == "example.com");
        UNIT_ASSERT(!converter.InitializeFromBaseUrl(""));
    }

    Y_UNIT_TEST(InitializeFromFormattedUrl) {
        TErasureUrlConverter converter;
        UNIT_ASSERT(converter.InitializeFromFormattedUrl("erasure_02+01_01_example.com"));
        UNIT_ASSERT(converter.GetBaseUrl() == "example.com");
        UNIT_ASSERT(converter.GetMainNumber() == 2);
        UNIT_ASSERT(converter.GetParityNumber() == 1);
        UNIT_ASSERT(converter.GetPartNumber() == 1);
        UNIT_ASSERT(!converter.InitializeFromFormattedUrl("erasure_02+1_01_example.com"));
        UNIT_ASSERT(!converter.InitializeFromFormattedUrl("erasure_02_01_01_example.com"));
        UNIT_ASSERT(!converter.InitializeFromFormattedUrl("erasure_02+01_01_"));
        UNIT_ASSERT(converter.InitializeFromFormattedUrl("erasure_23+11_17_example.com"));
        UNIT_ASSERT(converter.GetBaseUrl() == "example.com");
        UNIT_ASSERT(converter.GetMainNumber() == 23);
        UNIT_ASSERT(converter.GetParityNumber() == 11);
        UNIT_ASSERT(converter.GetPartNumber() == 17);
    }

    Y_UNIT_TEST(GetFormattedUrl) {
        TErasureUrlConverter converter;
        UNIT_ASSERT(converter.InitializeFromFormattedUrl("erasure_02+01_01_example.com"));
        UNIT_ASSERT_EQUAL(converter.GetBaseUrl(), "example.com");
        UNIT_ASSERT_EQUAL(converter.GetMainNumber(), 2);
        UNIT_ASSERT(converter.GetParityNumber() == 1);
        UNIT_ASSERT(converter.GetPartNumber() == 1);
        UNIT_ASSERT_EQUAL(converter.GetFormattedUrl(2), TString("erasure_02+01_02_example.com"));
        UNIT_ASSERT_EQUAL(converter.GetFormattedUrl(3), TString(""));
    }

    Y_UNIT_TEST(SetErasureParams) {
        TErasureUrlConverter converter;
        UNIT_ASSERT(!converter.SetErasureParams(-1, 3));
        UNIT_ASSERT(!converter.SetErasureParams(2, 101));
        UNIT_ASSERT(converter.SetErasureParams(2, 3));
        UNIT_ASSERT_EQUAL(converter.GetMainNumber(), 2);
        UNIT_ASSERT_EQUAL(converter.GetParityNumber(), 3);
    }
}

Y_UNIT_TEST_SUITE(TRtyShardingTests) {
    Y_UNIT_TEST(ShardingRestriction) {
        InitLog();
        NRTYServer::TMessage message;
        message.MutableDocument()->SetUrl("http://ПаНаМар.cOm/123456980_2");
        TShardsDispatcher::TContext context(NSaas::KeyPrefix, 5);
        TShardsDispatcher dispatcher(context);
        for (ui64 i = Max<ui64>(); i > (Max<ui64>() - (1 << 16)); --i) {
            // DEBUG_LOG << i << Endl;
            message.MutableDocument()->SetKeyPrefix(i);
            UNIT_ASSERT(dispatcher.GetShard(message) < NSearchMapParser::SearchMapShards);
        }
    }

    Y_UNIT_TEST(ShardingSimple) {
        InitLog();
        {
            TShardsDispatcher::TContext ctx(NSaas::KeyPrefix, 5);
            const TString serialized = ctx.ToString();
            TShardsDispatcher::TContext newCtx = TShardsDispatcher::TContext::FromString(serialized);
            const TString newSerialized = newCtx.ToString();
            UNIT_ASSERT_EQUAL(serialized, newSerialized);
            UNIT_ASSERT_EQUAL(ctx.Type, newCtx.Type);
            UNIT_ASSERT_EQUAL(ctx.KpsShift, newCtx.KpsShift);
        }

        {
            TShardsDispatcher::TContext ctx(NSaas::UrlHash);
            const TString serialized = ctx.ToString();
            Cout << serialized << Endl;
            TShardsDispatcher::TContext newCtx = TShardsDispatcher::TContext::FromString(serialized);
            const TString newSerialized = newCtx.ToString();
            UNIT_ASSERT_EQUAL(serialized, newSerialized);
            UNIT_ASSERT_EQUAL(ctx.Type, newCtx.Type);
            UNIT_ASSERT_EQUAL(ctx.KpsShift, newCtx.KpsShift);
        }
        {
            // Old format
            const TString serializedOld = "url-10";
            const TString serialized = "url_hash-10";
            TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString(serialized);
            TShardsDispatcher::TContext oldCtx = TShardsDispatcher::TContext::FromString(serializedOld);
            const TString newSerialized = oldCtx.ToString();
            UNIT_ASSERT_EQUAL(serialized, newSerialized);
            UNIT_ASSERT_EQUAL(ctx.Type, oldCtx.Type);
            UNIT_ASSERT_EQUAL(ctx.KpsShift, oldCtx.KpsShift);
        }
        {
            // Default
            const TString serialized = "broadcast";
            TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString(serialized);
            UNIT_ASSERT_EQUAL(ctx.Type, Broadcast);
            UNIT_ASSERT_EQUAL(ctx.KpsShift, 0);
        }
        {
            // Errors
            UNIT_ASSERT_EXCEPTION(TShardsDispatcher::TContext::FromString(""), yexception);
            UNIT_ASSERT_NO_EXCEPTION(TShardsDispatcher::TContext::FromString("url"));
            UNIT_ASSERT_EXCEPTION(TShardsDispatcher::TContext::FromString("1-1"), yexception);
            UNIT_ASSERT_EXCEPTION(TShardsDispatcher::TContext::FromString("kps-1-0"), yexception);
            UNIT_ASSERT_EXCEPTION(TShardsDispatcher::TContext::FromString("kps-17"), yexception);
            UNIT_ASSERT_EXCEPTION(TShardsDispatcher::TContext::FromString("key-17"), yexception);
        }
    }

    const TShardIntervals TestShardIntervals = {
        {0, 10921},
        {10922, 21483},
        {21844, 32765},
        {32766, 43687},
        {43688, 54609},
        {54610, 65533}
    };
    struct TTestShardIntervalEnumerator : IShardIntervalCallback {
        TSet<size_t> Intervals;

        void OnShardInterval(size_t index) override {
            Intervals.insert(index);
        }
    };

    void TestShardEnumerating(const TShardsDispatcher& dispatcher, TStringBuf url, ui64 kps) {
        TString strUrl{url};
        TSet<size_t> expected;
        for (size_t i = 0, imax = TestShardIntervals.size(); i < imax; ++i) {
            if (dispatcher.CheckSearchInterval(strUrl, kps, TestShardIntervals[i])) {
                expected.insert(i);
            }
        }
        TTestShardIntervalEnumerator enumerator;
        dispatcher.EnumerateIntervals(url, kps, TestShardIntervals, enumerator);
        UNIT_ASSERT_VALUES_EQUAL(expected, enumerator.Intervals);
    }

    Y_UNIT_TEST(GetUrlShardOnBoundaries) {
        const TShardIntervals intervals10 = {
            {0, 6522},
            {6523, 13105},
            {13106, 19658},
            {19659, 26211},
            {26212, 32764},
            {32765, 39317},
            {39318, 45870},
            {45871, 52423},
            {52424, 58976},
            {58977, 65533}
        };
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash");
        TShardsDispatcher dispatcher(ctx);
        TString strUrl{"sbp3537b305c11631b1afc79798e1f04a5575fb598f"};
        TSet<size_t> expected;
        for (size_t i = 0, imax = intervals10.size(); i < imax; ++i) {
            if (dispatcher.CheckSearchInterval(strUrl, 0, intervals10[i])) {
                expected.insert(i);
            }
        }
        TTestShardIntervalEnumerator enumerator;
        dispatcher.EnumerateIntervals(strUrl, 0, intervals10, enumerator);
        UNIT_ASSERT_VALUES_EQUAL(expected, enumerator.Intervals);
    }

    Y_UNIT_TEST(GetUrlShard) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash");
        TShardsDispatcher dispatcher(ctx);
        {
            NRTYServer::TMessage message;
            TString url = "url.com";
            message.MutableDocument()->SetUrl(url);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), ((ui32)CityHash64(url) % NSearchMapParser::SearchMapShards));
        }
        TestShardEnumerating(dispatcher, "url.com", 0);
        TestShardEnumerating(dispatcher, "something.ru", 100);
    }

    Y_UNIT_TEST(UrlShardingRule) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash");
        TUrlShardingRule rule(ctx);
        TVector<std::pair<TString, ui32>> url2shard {
                {"RLtE5Js6g58sFbBupH1b4ttHExbwAYsbslov2e1ik7Yzhk", 62946},
                {"Wzr73k1Kwwws6SkfITRBV6ABoIklzSdTRS", 62657},
                {"TAal6Whwc8eFO7haalkmZ0WIxwwRNM2", 11765},
                {"Xacd6HYgNCvkIaIsJ2eq189w43NQPNEw", 38274},
                {"oY146hRIb9YBRRVS5JTESNGVlgv6tp2hN3lTlBkMKIV", 34359},
                {"hyCnIdzk0DGjs9qlysSM", 43501},
                {"CneNV0fmUPTSZEiABTBPhsHog0m", 54709},
                {"S6Qes4zL3f5X2Zxalgkx7tKmmZ", 29895},
                {"QZKx1PNeHPbaSEgxHDvpyZZtgT", 12093},
                {"SDm0bJKYJjtxyJWXllDh0QapO", 43434},
                {"OInpTD0EDcXV8qD2l9LIg43tl3JLsgSEw", 57629},
                {"2PS0dda99jzKllT43A843s7LBx1c9xr9KIAytk70T4K", 60669},
                {"noGqyNszDzKMuLy2GoBb4J", 10764},
                {"x2cwVgGgU2uI0FByQ8In1cl3qXCRbWwXYYT1Dhf", 2372},
                {"kavkn4iEBz1CZkDPIPprtlOPI", 64134},
                {"PTX48Hc1zO6gcHfbRCLebb1OR", 39618},
                {"AnIQdhtAkzPK0Oi44KLhvkrFMToOpFFz3XPeD", 30208},
                {"ovHBFGZyKcGeHCO9PlaD7PRMNSrk641uxG3DWaBpCR", 15583},
                {"SbrZya7AhUb1r3ta9vb1tr46MeFBpeeFGToE4", 15012},
                {"MjNxlDycn7XO6QD8UXljiYOKncm", 12737},
                {"ooSBXn6G05snAO9HmLNHq8yZgKrTwd0jRSuxGyDE340", 19125},
                {"Q7TCSqTHxSqdmf7hH7QiYIfmGr0I", 33455},
                {"yUJ6vVWjxngPbK9HPqYUEE0Wsmlhl1ajTrnOKK5", 5314},
                {"7mgj4oyReUtriuNYe6FN7n5YeswxC", 57704},
                {"EIktzohxDJP6A7yWgdatygh3EJ", 2202},
                {"9G69TOtMNf1IsKhwTnvzTWZKUD1YSJSbZYk2vbMGGNZ6", 25770},
                {"G091v6SPdnuroQJ7qsfOafHB2YPP7T5LTENN", 64054},
                {"pmwaenP46Usw7HWnWhp44E", 41316},
                {"XKV0WqLpd8lEmYbpc5hYBw2yTjnXlzhIHAGo0abbiMpSsQH", 35112},
                {"Ux0ft0dKhOsSNXifhwRfW2GCO4fEUugwreZleA3LYv", 39290},
                {"uQTZ5P0kw2z8O3lcVgIQXwzGY8KtgDgYe77jW5SQ5PXRSI3vy", 32118},
                {"LUDJilR2e5EIesNjZupzKsOpKq8Lk4d5WqMFalHCrVT", 39598},
                {"LqCvkzs3sHsA5yTp16tXvFmT1", 59933},
                {"fqX6JtxVmfVegLtXTzVwoU1h2uvo", 37771},
                {"uq0kO73F2WzhzbMIVtlUMgG", 36179},
                {"fwtZpfujWuRszU71y4hyf3G8u0aGer7JlyqaeiHAai2aA", 43723},
                {"Z8BG6oHLxBLXQxOVG7U4h6mwEMeEwoLWxUk", 56035},
                {"I2KFBdnba97qHbswfettzX7vKQPHu", 7150},
                {"GC8bQJEBvpK2dZbWUG8N77sE1AcOPWWggeF4yTFrGYts7S", 24800},
                {"bgUwncnzbvZzL3WzhAEm6f1xwZoofEAEucBfEw", 17953},
                {"prBpCEJBwRNgY3fTxFGjsUt6NTH0iB", 35307},
                {"V1mIBzZLt1iZXlE0iRGPJjI", 23946},
                {"6mf6SphvqUC1RlMKm5rHo5IVVWumFBALvqSLd7ETbQ", 55198},
                {"2AnKuqBCeFS7ky0UE9fZ4VzN67bXg0zh8m12cD", 61562},
                {"FSeKAdKerRJOVmMqsRQP7Q", 18154},
                {"mY7oYkzlxRPFZ2YBrzVGTFaJ61j", 26144},
                {"67qtdyfCgevb3KqcLOlCOGQpU08a1", 50372},
                {"36vuzZqcjWGc7HwxtHKeSgts", 19573},
                {"NtdvsRyyNqvugXbClDJ27FtNZWpFPhkkaLdRlbP6S", 27356},
                {"06qZJamacrnWEM03ZNlH7JbikMJ8SjRQqFz7p", 1917}
        };
        for (auto&& [url, shard] : url2shard) {
            UNIT_ASSERT_VALUES_EQUAL(rule.GetUrlShard(url, NSearchMapParser::SearchMapShards), shard);
        }
    }

    Y_UNIT_TEST(GetUrlToLastOctothorp) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_to_last_octothorp");
        TShardsDispatcher dispatcher(ctx);
        for (auto&& prefix : {"url.com", "url.com#next", "x#y#z"}) {
            const ui32 shard = ((ui32)CityHash64(TStringBuf{prefix}) % NSearchMapParser::SearchMapShards);
            for (auto&& suffix : {"", "#", "#1", "#abcd"}) {
                if (TStringBuf{suffix} == "" && TStringBuf{prefix}.Contains('#')) {
                    // Not applicable because the prefix itself gets split.
                    // "url.com#next" is shareded as "url.com" but "url.com#next#abcd" is sharded as "url.com#next"
                    // That is the expected behavior for the "to the last" rule.
                    continue;
                }
                const auto url = TString::Join(prefix, suffix);
                NRTYServer::TMessage message;
                message.MutableDocument()->SetUrl(url);
                UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), shard);

                TestShardEnumerating(dispatcher, url, 0);
                TestShardEnumerating(dispatcher, url, 100);
            }
        }
    }

    Y_UNIT_TEST(GetUrlToLastUnderscore) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_to_last_underscore");
        TShardsDispatcher dispatcher(ctx);
        for (auto&& prefix : {"url.com", "url.com_next", "x_y_z"}) {
            const ui32 shard = ((ui32)CityHash64(TStringBuf{prefix}) % NSearchMapParser::SearchMapShards);
            for (auto&& suffix : {"", "_", "_1", "_abcd"}) {
                if (TStringBuf{suffix} == "" && TStringBuf{prefix}.Contains('_')) {
                    // Not applicable because the prefix itself gets split.
                    // "url.com_next" is shareded as "url.com" but "url.com_next_abcd" is sharded as "url.com_next"
                    // That is the expected behavior for the "to the last" rule.
                    continue;
                }
                const auto url = TString::Join(prefix, suffix);
                NRTYServer::TMessage message;
                message.MutableDocument()->SetUrl(url);
                UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), shard);

                TestShardEnumerating(dispatcher, url, 0);
                TestShardEnumerating(dispatcher, url, 100);
            }
        }
    }

    Y_UNIT_TEST(ErasureFixed) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash_erasure");
        TShardsDispatcher dispatcher(ctx);
        NRTYServer::TMessage message;
        TString url = "erasure_03+02_03_you_shall_not_pass";
        message.MutableDocument()->SetUrl(url);
        NSearchMapParser::TShardIndex ansShardIndex = 19774;
        UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), ansShardIndex);
    }

    Y_UNIT_TEST(ErasureEnumerateIntervals) {
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash_erasure");
        TShardsDispatcher dispatcher(ctx);
        TString url1{"erasure_02+01_00_url"};
        TString url2{"erasure_02+01_02_test_url"};
        TestShardEnumerating(dispatcher, url1, 0);
        TestShardEnumerating(dispatcher, url2, 0);
    }

    Y_UNIT_TEST(ErasureShardingRestriction) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash_erasure");
        TShardsDispatcher dispatcher(ctx);
        NRTYServer::TMessage message;
        // Computing the shard index
        TString url = "erasure_02+01_00_url.com";
        message.MutableDocument()->SetUrl(url);
        auto shard = dispatcher.GetShard(message);
        auto step = dispatcher.GetContext().GetShardsMax() / 3;
        // Checking the boundaries
        UNIT_ASSERT(shard >= 0);
        UNIT_ASSERT(shard < step * 3);
    }

    Y_UNIT_TEST(GetUrlErasureShardStep) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("url_hash_erasure");
        TShardsDispatcher dispatcher(ctx);
        // Computing shard indexes for the erasure-formatted parts of url
        NRTYServer::TMessage message;
        TString url = "erasure_02+01_00_url.com";
        message.MutableDocument()->SetUrl(url);
        auto shard = dispatcher.GetShard(message);
        NRTYServer::TMessage message1;
        TString url1 = "erasure_02+01_01_url.com";
        message1.MutableDocument()->SetUrl(url1);
        auto shard1 = dispatcher.GetShard(message1);
        NRTYServer::TMessage message2;
        TString url2 = "erasure_02+01_02_url.com";
        message2.MutableDocument()->SetUrl(url2);
        auto shard2 = dispatcher.GetShard(message2);
        auto step = dispatcher.GetContext().GetShardsMax() / 3;
        // Checking the distance to be a multiple of step
        UNIT_ASSERT((shard2 - shard1 + step * 3) % step == 0);
        UNIT_ASSERT((shard2 - shard + step * 3) % step == 0);
        UNIT_ASSERT((shard - shard1 + step * 3) % step == 0);
        UNIT_ASSERT(shard1 != shard);
        UNIT_ASSERT(shard != shard2);
        UNIT_ASSERT(shard1 != shard2);
    }

    Y_UNIT_TEST(GetKeyPrefixShard) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("keyprefix");
        TShardsDispatcher dispatcher(ctx);
        TShardsDispatcher::TContext ctx1 = TShardsDispatcher::TContext::FromString("keyprefix-1");
        TShardsDispatcher dispatcher1(ctx1);
        TString url = "url.com";
        ui32 urlShard = ((ui32)CityHash64(url) % NSearchMapParser::SearchMapShards);
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetKeyPrefix(0);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 0);
        }
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetUrl(url);
            message.MutableDocument()->SetKeyPrefix(0);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 0);
        }
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetUrl(url);
            message.MutableDocument()->SetKeyPrefix(5);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 5);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher1.GetShard(message) % 2, urlShard % 2);
        }
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetUrl(url);
            message.MutableDocument()->SetKeyPrefix(NSearchMapParser::SearchMapShards + 5);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 5);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher1.GetShard(message) % 2, urlShard % 2);
        }
        TestShardEnumerating(dispatcher, "url.com", 0);
        TestShardEnumerating(dispatcher, "something.ru", 100);
        TestShardEnumerating(dispatcher1, "url.com", 0);
        TestShardEnumerating(dispatcher1, "something.ru", 100);
    }

    Y_UNIT_TEST(GetBroadcastShard) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("broadcast");
        TShardsDispatcher dispatcher(ctx);
        TString url = "url.com";
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetKeyPrefix(0);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 0);
        }
        {
            NRTYServer::TMessage message;
            message.MutableDocument()->SetUrl(url);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 0);
        }
        TestShardEnumerating(dispatcher, "url.com", 0);
        TestShardEnumerating(dispatcher, "something.ru", 100);
    }

    Y_UNIT_TEST(GetExternalShard) {
        InitLog();
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("external");
        TShardsDispatcher dispatcher(ctx);
        {
            NRTYServer::TMessage message;
            message.SetExternalShard(123);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 123);
        }
        {
            NRTYServer::TMessage message;
            message.SetExternalShard(5 * NSearchMapParser::SearchMapShards + 123);
            UNIT_ASSERT_VALUES_EQUAL(dispatcher.GetShard(message), 123);
        }
        TestShardEnumerating(dispatcher, "url.com", 0);
        TestShardEnumerating(dispatcher, "something.ru", 100);
    }

    Y_UNIT_TEST(QuerySearch) {
        InitLog();
        const TShardsDispatcher::TContext ctxUrlHash = TShardsDispatcher::TContext::FromString("url_hash");
        const TShardsDispatcher::TContext ctxQuerySearch = TShardsDispatcher::TContext::FromString("querysearch");
        const TShardsDispatcher urlHash(ctxUrlHash);
        const TShardsDispatcher querySearch(ctxQuerySearch);
        {
            auto checkUrls = [&](TStringBuf qsKey, TStringBuf uhKey) {
                auto qsShard = querySearch.GetShard(qsKey, 0);
                auto uhShard = urlHash.GetShard(uhKey, 0);
                UNIT_ASSERT_VALUES_EQUAL(qsShard, uhShard);
            };
            checkUrls(".7\thttps://yandex.ru/", ".7\thttps://yandex.ru/");
            checkUrls(".7.5\thttps://yandex.ru/\t225", ".7.5\thttps://yandex.ru/");
            checkUrls(".9\texample.com", ".9\texample.com");
            checkUrls(".9\texample.com/some/path?x=10", ".9\texample.com");
            checkUrls(".99\texample.com/something.html", ".99\texample.com/something.html");
        }
        TestShardEnumerating(querySearch, ".7\thttps://yandex.ru/", 0);
        TestShardEnumerating(querySearch, ".9\texample.com/", 100);
    }

    void CheckSearchInterval(const TShardsDispatcher& dispatcher, ui64 kps) {
        NRTYServer::TMessage message;
        message.MutableDocument()->SetKeyPrefix(kps);
        message.MutableDocument()->SetUrl("123");
        ui64 shard = dispatcher.GetShard(message);
        ui64 minShard = shard >> dispatcher.GetContext().KpsShift << dispatcher.GetContext().KpsShift;
        ui64 maxShard = minShard + (1 << dispatcher.GetContext().KpsShift) - 1;
        NSearchMapParser::TShardsInterval interval(minShard, maxShard);
        UNIT_ASSERT(dispatcher.CheckSearchInterval(TString(), kps, interval));
    }

    Y_UNIT_TEST(ShardingCheckInterval) {
        InitLog();
        {
            TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("keyprefix");
            TShardsDispatcher dispatcher(ctx);
            NSearchMapParser::TShardsInterval interval(0, 123);
            UNIT_ASSERT(dispatcher.CheckSearchInterval(TString(), 120, interval));
            UNIT_ASSERT(dispatcher.CheckSearchInterval(TString(), 123, interval));
            UNIT_ASSERT(!dispatcher.CheckSearchInterval(TString(), 124, interval));
            UNIT_ASSERT(!dispatcher.CheckSearchInterval(TString(), 1000, interval));
        }
        ///
        {
            TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("keyprefix-3");
            TShardsDispatcher dispatcher(ctx);
            CheckSearchInterval(dispatcher, 32);
            CheckSearchInterval(dispatcher, 33);
            CheckSearchInterval(dispatcher, 39);
            CheckSearchInterval(dispatcher, 30);
            CheckSearchInterval(dispatcher, 24);
            CheckSearchInterval(dispatcher, 40);
        }
    }
}

Y_UNIT_TEST_SUITE(TRtyShardingRaspr) {
    Y_UNIT_TEST(GetKeyPrefixShard) {
        return;
        TShardsDispatcher::TContext ctx = TShardsDispatcher::TContext::FromString("keyprefix-11");
        TShardsDispatcher dispatcher(ctx);
        ui32 intervalSize = NSearchMapParser::SearchMapShards / 300;
        TVector<ui32> raspr(NSearchMapParser::SearchMapShards / intervalSize + 1, 0);
        ui64 kps = 1;
        for (ui32 i = 0; i < 100000; ++i) {
            TString url = "url.com/";
            NRTYServer::TMessage message;
            message.MutableDocument()->SetUrl(url);
            message.MutableDocument()->SetKeyPrefix(kps++);
            raspr[dispatcher.GetShard(message) / intervalSize]++;
        }
        Sort(raspr.begin(), raspr.end());
        TMap<i32, ui32> freq;
        for (ui32 count: raspr) {
            Cout << count << Endl;
            freq[count]++;
        }
        Cout << raspr.front() << "-" << raspr.back() << Endl;
        i32 prev = -1;
        for (const auto& f : freq) {
            if (prev > 0)
                for(;prev < f.first;++prev)
                    Cout << Endl;
            else
                prev = f.first;
            Cout << TString(f.second, '*') << Endl;
        }
    }
}

static void CheckEqual(TStringBuf actual, TStringBuf expected1, TStringBuf expected2) {
    if (expected1.empty()) {
        UNIT_ASSERT_EQUAL(actual, expected2);
    } else {
        UNIT_ASSERT_EQUAL(actual, expected1);
    }
}

template<typename TEnumType>
void TestEnumSerialization(const TString& serializedEnum, TStringBuf expectedString = {}) {
    auto parsedValue = FromString<TEnumType>(serializedEnum);
    TString actualString = ToString(parsedValue);
    CheckEqual(actualString, expectedString, serializedEnum);

    TString buffer;
    {
        TStringOutput strStream(buffer);
        strStream << parsedValue;
    }
    CheckEqual(buffer, expectedString, serializedEnum);
}

void TestEnumProtoConversation(NSaas::ShardingType sharding, NSaas::ShardsCount shardsCount) {
    NSaasProto::TService service;
    {
        TShardsDispatcher::TContext context(sharding);
        context.Shards = shardsCount;
        context.ToProto(service);
    }
    TShardsDispatcher::TContext context = TShardsDispatcher::TContext::FromProto(service);
    UNIT_ASSERT_EQUAL(sharding, context.Type);
    UNIT_ASSERT_EQUAL(shardsCount, context.Shards);
}

Y_UNIT_TEST_SUITE(SaasShardingTypeSuite) {
    Y_UNIT_TEST(ShardingTypeSerialization) {
        for (auto enumItem :  GetEnumNames<NSaas::ShardingType>()) {
            TestEnumSerialization<NSaas::ShardingType>(enumItem.second);
        }
        TestEnumSerialization<NSaas::ShardingType>("key", "keyprefix");
        TestEnumSerialization<NSaas::ShardingType>("url", "url_hash");
        UNIT_ASSERT_EXCEPTION(FromString<NSaas::ShardingType>("something wrong"), yexception);
    }
    Y_UNIT_TEST(ShardsCountSerialization) {
        for (auto enumItem :  GetEnumNames<NSaas::ShardsCount>()) {
            TestEnumSerialization<NSaas::ShardsCount>(enumItem.second);
        }
        UNIT_ASSERT_EXCEPTION(FromString<NSaas::ShardsCount>("something wrong"), yexception);
    }
    Y_UNIT_TEST(ShardingTypeProtoConversation) {
        for (auto value : GetEnumAllValues<NSaas::ShardingType>()) {
            TestEnumProtoConversation(value, ShardsCount::Legacy);
        }
    }
    Y_UNIT_TEST(ShardsCountProtoConversation) {
        for (auto value : GetEnumAllValues<NSaas::ShardsCount>()) {
            TestEnumProtoConversation(UrlHash, value);
        }
    }
}

