#include "donors_builder.h"

#include <saas/library/sharding/sharding.h>

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

using namespace NSearchMapParser;

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

    TVector<TSearchMapHost> GenerateSharding(const ui32 shardsCount, const TString& hostPrefix) {
        TVector<TSearchMapHost> hosts;
        hosts.reserve(shardsCount);
        for (ui32 i = 0; i < shardsCount; ++i) {
            const TShardsInterval interval = NSaas::TSharding::GetInterval(i, shardsCount);
            const TString hostName = hostPrefix + ToString(i);
            hosts.emplace_back(hostName, interval, Master, 1000 + i, 1002 + i);
        }
        return hosts;
    }

    TShardsInterval MergeSolid(const TVector<NRTYReplication::TIntervalLink>& intervalLinks) {
        UNIT_ASSERT(!intervalLinks.empty());
        TVector<TShardsInterval> intervals;
        intervals.reserve(intervalLinks.size());
        for(const NRTYReplication::TIntervalLink& intervalLink: intervalLinks) {
            intervals.push_back(intervalLink.Donor->Shards);
        }
        Sort(intervals);
        TShardsInterval res(intervals.front());
        for (ui32 i = 1; i < intervals.size(); ++i) {
            const TShardsInterval& interval = intervals[i];
            UNIT_ASSERT_EQUAL_C(res.GetMax() + 1, interval.GetMin(), ToString(res.GetMax() + 1) + "==" + ToString(interval.GetMin()));
            res.SetMax(interval.GetMax());
        }
        return res;
    }

    void TestReshard(const ui32 fromShardsCount, const ui32 toShardsCount) {
        InitLog();
        TVector<TSearchMapHost> shardingFrom = GenerateSharding(fromShardsCount, "src");
        TVector<TSearchMapHost> shardingTo = GenerateSharding(toShardsCount, "dst");
        NRTYReplication::TDonorsBuilder donorsBuilder(shardingTo, shardingFrom);
        UNIT_ASSERT(donorsBuilder.BuildRemap());
        for (const TSearchMapHost& dstSlot: shardingTo) {
            TVector<NRTYReplication::TIntervalLink> restoreLink = donorsBuilder.GetRestoreInfo(dstSlot);
            UNIT_ASSERT_C(!restoreLink.empty(), "No interval for " + dstSlot.GetSlotName() + ":" + dstSlot.Shards.ToString());
            TShardsInterval merged = MergeSolid(restoreLink);
            UNIT_ASSERT(dstSlot.Shards.GetMin() <= dstSlot.Shards.GetMax());
            UNIT_ASSERT_C(merged.GetMin() <= dstSlot.Shards.GetMin(), ToString(merged.GetMin()) + "<=" + ToString(dstSlot.Shards.GetMin()));
            UNIT_ASSERT_C(merged.GetMax() >= dstSlot.Shards.GetMax(), ToString(merged.GetMax()) + ">=" + ToString(dstSlot.Shards.GetMax()));
        }
    }
}


Y_UNIT_TEST_SUITE(TDonorsBuilderTest) {

    Y_UNIT_TEST(BuildDonors1to3) {
        TestReshard(1, 3);
    };

    Y_UNIT_TEST(BuildDonors3to1) {
        TestReshard(3, 1);
    };

    Y_UNIT_TEST(BuildDonors1to4) {
        TestReshard(1, 4);
    };

    Y_UNIT_TEST(BuildDonors4to1) {
        TestReshard(4, 1);
    };

    Y_UNIT_TEST(BuildDonors3to5) {
        TestReshard(3, 5);
    };

    Y_UNIT_TEST(BuildDonors5to3) {
        TestReshard(5, 3);
    };

    Y_UNIT_TEST(BuildDonors1to7) {
        TestReshard(1, 7);
    };

    Y_UNIT_TEST(BuildDonors7to1) {
        TestReshard(7, 1);
    };

    Y_UNIT_TEST(BuildDonors2to10) {
        TestReshard(2, 10);
    };

    Y_UNIT_TEST(BuildDonors10to2) {
        TestReshard(10, 2);
    };

    Y_UNIT_TEST(BuildDonors10to40) {
        TestReshard(10, 40);
    };

    Y_UNIT_TEST(BuildDonors40to10) {
        TestReshard(40, 10);
    };

}
