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


using namespace NRTYDeploy;

Y_UNIT_TEST_SUITE(TRtySlotsPlacementTest) {

    void AddSlot(TMap<TString, NRTYCluster::TSlotData>& slots, NRTYCluster::TSlotData slot) {
        TString name = slot.ShortSlotName();
        slots.emplace(std::move(name), std::move(slot));
    }

    Y_UNIT_TEST(SlotsPlacementTrivial1) {
        TMap<TString, NRTYCluster::TSlotData> slots;
        const NRTYCluster::TSlotData slot("man1-0001.search.yandex.net", 1000);
        AddSlot(slots, slot);
        TMap<NSearchMapParser::TShardsInterval, THashSet<TString>> forbiddenSet;
        TVector<NSearchMapParser::TShardsInterval> intervals;
        intervals.emplace_back(NSearchMapParser::FullInterval);

        TMaybe<TVector<NRTYCluster::TSlotData>> res = PlaceSlots(slots, forbiddenSet, intervals);
        UNIT_ASSERT(res.Defined());
        UNIT_ASSERT_EQUAL(res.GetRef().size(), 1);
        UNIT_ASSERT_EQUAL(res.GetRef()[0], slot);
    }

    Y_UNIT_TEST(SlotsPlacementTrivial2) {
        TMap<TString, NRTYCluster::TSlotData> slots;
        const NRTYCluster::TSlotData slot1("man1-0001.search.yandex.net", 1000);
        const NRTYCluster::TSlotData slot2("man1-0001.search.yandex.net", 2000);
        AddSlot(slots, slot1);
        AddSlot(slots, slot2);

        TMap<NSearchMapParser::TShardsInterval, THashSet<TString>> forbiddenSet;
        TVector<NSearchMapParser::TShardsInterval> intervals = NSaas::TSharding::SplitInterval(NSearchMapParser::FullInterval, 2);
        TMaybe<TVector<NRTYCluster::TSlotData>> res = PlaceSlots(slots, forbiddenSet, intervals);
        UNIT_ASSERT(res.Defined());
        UNIT_ASSERT_EQUAL(res.GetRef().size(), 2);
    }

    Y_UNIT_TEST(SlotsPlacementForbidden2) {
        TMap<TString, NRTYCluster::TSlotData> slots;
        const NRTYCluster::TSlotData slot1("man1-0001.search.yandex.net", 1000);
        const NRTYCluster::TSlotData slot2("man1-0002.search.yandex.net", 1000);
        AddSlot(slots, slot1);
        AddSlot(slots, slot2);

        TMap<NSearchMapParser::TShardsInterval, THashSet<TString>> forbiddenSet;
        TVector<NSearchMapParser::TShardsInterval> intervals = NSaas::TSharding::SplitInterval(NSearchMapParser::FullInterval, 2);
        forbiddenSet.insert(std::make_pair(intervals[0], THashSet<TString>({slot1.ShortHost()})));
        forbiddenSet.insert(std::make_pair(intervals[1], THashSet<TString>({slot2.ShortHost()})));
        TMaybe<TVector<NRTYCluster::TSlotData>> res = PlaceSlots(slots, forbiddenSet, intervals);
        UNIT_ASSERT(res.Defined());
        UNIT_ASSERT_EQUAL(res.GetRef().size(), 2);
        UNIT_ASSERT_EQUAL(res.GetRef()[0], slot2);
        UNIT_ASSERT_EQUAL(res.GetRef()[1], slot1);
    }

    Y_UNIT_TEST(SlotsPlacementForbiddenImpossible) {
        TMap<TString, NRTYCluster::TSlotData> slots;
        const NRTYCluster::TSlotData slot1("man1-0001.search.yandex.net", 1000);
        const NRTYCluster::TSlotData slot2("man1-0002.search.yandex.net", 1000);
        AddSlot(slots, slot1);
        AddSlot(slots, slot2);

        TMap<NSearchMapParser::TShardsInterval, THashSet<TString>> forbiddenSet;
        TVector<NSearchMapParser::TShardsInterval> intervals = NSaas::TSharding::SplitInterval(NSearchMapParser::FullInterval, 2);
        forbiddenSet.insert(std::make_pair(intervals[0], THashSet<TString>({slot1.ShortHost()})));
        forbiddenSet.insert(std::make_pair(intervals[1], THashSet<TString>({slot1.ShortHost()})));
        TMaybe<TVector<NRTYCluster::TSlotData>> res = PlaceSlots(slots, forbiddenSet, intervals);
        UNIT_ASSERT(!res.Defined());
    }

}
