#include "task.h"
#include <saas/deploy_manager/scripts/add_slots/action.h>
#include <saas/deploy_manager/scripts/allocate_slots/action/action_same.h>
#include <saas/deploy_manager/scripts/allocate_slots/common/intervals.h>
#include <saas/deploy_manager/scripts/common/dead_slots_collector/dead_slots_collector.h>
#include <saas/deploy_manager/scripts/common/deploy/deploy_builder.h>
#include <saas/deploy_manager/scripts/release_slots/action.h>
#include <saas/deploy_manager/scripts/restore_index/action.h>
#include <saas/deploy_manager/scripts/searchmap/action.h>
#include <util/generic/utility.h>

namespace NRTYDeploy {

    TReshardTask::TReshardTask(
        NSearchMapParser::TShardIndex shardMin, NSearchMapParser::TShardIndex shardMax, ui32 newPartsCount,
        const TClusterTask::TCgiContext& context, NRTYDeploy::ICommonData* commonData)
        : TClusterTask(context, commonData, RESHARD_TASK_TYPE)
        , ShardMin(shardMin)
        , ShardMax(shardMax)
        , NewPartsCount(newPartsCount)
    {
    }

    void TReshardTask::DoBuildTask() {

        TVector<TString> oldSlots;

        NRTYDeployInfo::IDeployComponentInfo::TPtr info = NRTYDeployInfo::IDeployComponentInfo::TFactory::Construct(ServiceType);
        if (!info) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "Incorrect service type: " + ServiceType);
            return;
        }

        info->SetInfo(CommonData, CType);
        auto ism = info->SearchMap(Service);

        if (ism.GetInternalSearchMap().size() != 1) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "incorrect searchmap for " + CType);
            return;
        }
        NSearchMapParser::TSearchMapService* serviceObj = ism.GetService(Service);
        if (!serviceObj) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "Service " + Service + " doesn't exists");
            return;
        }

        NSearchMapParser::TSearchMapReplica* replica = serviceObj->GetReplica("default");
        if (!replica) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "ConfigType default doesn't exists");
            return;
        }

        TVector<NSearchMapParser::TSearchMapHost> slots = replica->GetSlotsContainedIn(ShardMin, ShardMax);
        TMap<TString, TMap<NSearchMapParser::TShardsInterval, TVector<NSearchMapParser::TSearchMapHost>>> hostsByDC;
        for (ui32 sl = 0; sl < slots.size(); ++sl) {
            NRTYCluster::TSlotData sd(slots[sl].Name, slots[sl].SearchPort);
            hostsByDC[sd.GetDC()][slots[sl].Shards].push_back(slots[sl]);
            oldSlots.push_back(slots[sl].GetSlotName());
        }

        TMap<TString, ui32> replicasByDC;
        for (auto&& i : hostsByDC) {
            for (auto&& r : i.second) {
                replicasByDC[i.first] = Max<ui32>(replicasByDC[i.first], r.second.size());
            }
        }

        NDaemonController::TIntervalsByDC intervalsByDC;
        for (auto&& i : replicasByDC) {
            for (ui32 s = 0; s < NewPartsCount; ++s) {
                for (ui32 r = 0; r < i.second; ++r) {
                    NSearchMapParser::TShardsInterval interval = NSaas::TSharding::GetInterval(s, NewPartsCount, ShardMin, ShardMax);
                    intervalsByDC.AddInfo(i.first, interval);
                }
            }
        }


        NDaemonController::TAction::TPtr actionAllocate = new NDaemonController::TAllocateSameSlotsAction(CType, Service, RTYSERVER_SERVICE, intervalsByDC, TVector<TString>());
        NRTYScript::TSlotTaskContainer allocateNode = GetScript()->AddAction(CommonData->GetDeployManagerBalanserHost(), CommonData->GetDeployManagerBalanserPort(), CommonData->GetDeployManagerBalanserUriPrefix(), actionAllocate);

        TString slotsPoolAddress = "$" + allocateNode.GetName() + "-action.task.result_pool";

        NDaemonController::TAction::TPtr actionAddSlots = new NDaemonController::TActivateSlotsAction(slotsPoolAddress, CType, Service, ServiceType);
        NRTYScript::TSlotTaskContainer currentNode = allocateNode.AddNext(*GetScript(), actionAddSlots);

        NDaemonController::TAction::TPtr action = new NRTYDeploy::TRestoreIndexAction(Service, CType, TDeadSlotsCollector::GetNotUsefulSlots(Service, CommonData, ism),
            "$" + allocateNode.GetName() + "-action.task.result_pool", NDaemonController::apStartAndWait);
        currentNode = currentNode.AddNext(*GetScript(), action);

        NDaemonController::TAction::TPtr actionActivateSlots = new NDaemonController::TSimpleSearchmapModifAction(slotsPoolAddress, CType, Service, "enable_search,enable_indexing", ServiceType, true, false);
        currentNode = currentNode.AddNext(*GetScript(), actionActivateSlots);

        NDaemonController::TAction::TPtr actionRelease = new NDaemonController::TReleaseSlotsAction(oldSlots, CType, Service, ServiceType, true, true);
        currentNode = currentNode.AddNext(*GetScript(), actionRelease);

    }

};

