#include "allocator.h"
#include <saas/util/cluster/cluster.h>
#include <library/cpp/testing/unittest/registar.h>


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

Y_UNIT_TEST_SUITE(TRtySlotsAllocatorTest) {
    Y_UNIT_TEST(SlotsRegistration) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(!cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
    }

    Y_UNIT_TEST(SlotsCustomFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h2", 7200)));

        NRTYCluster::TSlotsAllocator allocator;
        TSet<NRTYCluster::TSlotData> slots;
        slots.insert(NRTYCluster::TSlotData("dc1__h2", 7100));
        slots.insert(NRTYCluster::TSlotData("dc1__h2", 7200));
        NRTYCluster::TCustomSlotsFilter filter(slots);

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));

        NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

        UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 2);
        UNIT_ASSERT(filtered->GetSlots().contains("dc1__h2:7100"));
        UNIT_ASSERT(filtered->GetSlots().contains("dc1__h2:7200"));

    };

    Y_UNIT_TEST(SlotsDCResourcesFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h3", 7000)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));

        {
            NRTYCluster::TDCResourcesFilter filter(3, false);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 6);
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7100"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7200"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }

        {
            NRTYCluster::TDCResourcesFilter filter(4, true);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 0);
        }

        {
            NRTYCluster::TDCResourcesFilter filter(3, true);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 3);
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }

        {
            NRTYCluster::TDCResourcesFilter filter(2, false);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 8);
        }

    };


    Y_UNIT_TEST(SlotsDCPreferredFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h3", 7000)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));

        {
            NRTYCluster::TPreferredDCListFilter filter({"dc1"});

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 2);
            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7100"));
        }

        {
            NRTYCluster::TPreferredDCListFilter filter({"dc1", "dc3"});

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 5);
            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7100"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }

        {
            NRTYCluster::TPreferredDCListFilter filter({});

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 0);
        }

        {
            NRTYCluster::TPreferredDCListFilter filter({"dc2", "dc3"});

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 6);

            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7100"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7200"));

            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }

    };

    Y_UNIT_TEST(SlotPerServerFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h3", 7000)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));

        {
            NRTYCluster::TOneSlotFilter filter;

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 5);

            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7000"));

            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }

        {
            NRTYCluster::TUsedServersFilter filter;

            filter.AddInPool("dc1__h1");
            filter.AddInPool("dc2__h2");
            filter.AddInPool("dc3__h1");

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(filter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 2);

            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }
    };

    Y_UNIT_TEST(SlotPoolFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc1__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc3__h3", 7000)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));

        {
            NRTYCluster::TPoolFilter pFilter(true);
            pFilter.AddInPool("dc1__h1:7000");
            pFilter.AddInPool("dc2__h2:7000");
            pFilter.AddInPool("dc2__h2:7200");
            pFilter.AddInPool("dc3__h1:7000");
            pFilter.AddInPool("dc3__h3:7000");

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(pFilter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 3);

            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7100"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7100"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h2:7000"));
        }

        {
            NRTYCluster::TPoolFilter pFilter(false);
            pFilter.AddInPool("dc1__h1:7000");
            pFilter.AddInPool("dc2__h2:7000");
            pFilter.AddInPool("dc2__h2:7200");
            pFilter.AddInPool("dc3__h1:7000");
            pFilter.AddInPool("dc3__h3:7000");

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(pFilter);

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 5);

            UNIT_ASSERT(filtered->GetSlots().contains("dc1__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7200"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc3__h3:7000"));
        }
    };

    Y_UNIT_TEST(SlotCorrectDistributionFilter) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h2", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc4__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc4__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc4__h3", 7000)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));
        NRTYCluster::TSlotsAllocator allocator;

        {
            TMap<TString, ui32> usedCount;
            TMap<TString, ui32> currentUsedCount;

            NRTYCluster::TCommonRankerFilter rFilter(usedCount, currentUsedCount, allocator, 3);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(rFilter);

            for (auto&& i : filtered->GetSlots()) {
                DEBUG_LOG << i.second.FullSlotName() << Endl;
            }

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 3);

            UNIT_ASSERT(filtered->GetSlots().contains("dc4__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc4__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc4__h3:7000"));
        }

        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h3", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc2__h4", 7000)));

        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));
        {
            TMap<TString, ui32> usedCount;
            TMap<TString, ui32> currentUsedCount;

            NRTYCluster::TCommonRankerFilter rFilter(usedCount, currentUsedCount, allocator, 3);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(rFilter);

            for (auto&& i : filtered->GetSlots()) {
                DEBUG_LOG << i.second.FullSlotName() << Endl;
            }

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 3);

            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc2__h3:7000"));
        }

    };

    Y_UNIT_TEST(SlotCorrectDistributionDiffServers) {
        InitLog();
        NRTYCluster::TCluster cluster;
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h1", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h1", 7100)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h1", 7200)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h2", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h3", 7000)));
        UNIT_ASSERT(cluster.RegisterSlot("1", NRTYCluster::TSlotData("dc0__h3", 7100)));

        NRTYCluster::TCTypeCluster cTypeCluster;
        UNIT_ASSERT(cluster.GetCTypeCluster("1", cTypeCluster));
        NRTYCluster::TSlotsAllocator allocator;

        {
            TMap<TString, ui32> usedCount;
            TMap<TString, ui32> currentUsedCount;

            NRTYCluster::TCommonRankerFilter rFilter(usedCount, currentUsedCount, allocator, 3);

            NRTYCluster::TCTypeCluster::TPtr filtered = cTypeCluster.FilterSlots(rFilter);

            for (auto&& i : filtered->GetSlots()) {
                DEBUG_LOG << i.second.FullSlotName() << Endl;
            }

            UNIT_ASSERT_EQUAL(filtered->GetSlots().size(), 3);

            UNIT_ASSERT(filtered->GetSlots().contains("dc0__h1:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc0__h2:7000"));
            UNIT_ASSERT(filtered->GetSlots().contains("dc0__h3:7000"));
        }

    };
}
