#include "task.h"
#include "donors_builder.h"
#include <saas/deploy_manager/meta/cluster.h>
#include <saas/deploy_manager/scripts/add_slots/action.h>
#include <saas/deploy_manager/scripts/common/deploy/deploy_builder.h>
#include <saas/deploy_manager/scripts/common/scripts_helper.h>
#include <saas/deploy_manager/scripts/deploy/action.h>
#include <saas/deploy_manager/scripts/logic/action.h>
#include <saas/deploy_manager/scripts/release_slots/action.h>
#include <saas/deploy_manager/scripts/searchmap/action.h>
#include <saas/rtyserver/controller/controller_actions/clear_index_action.h>
#include <saas/rtyserver/controller/controller_actions/detach_action.h>
#include <saas/rtyserver/controller/controller_actions/download_action.h>
#include <saas/rtyserver/controller/controller_actions/shards_action.h>
#include <saas/library/daemon_base/controller_actions/restart.h>
#include <saas/library/searchmap/searchmap.h>

namespace NRTYDeploy {

    class TScannerInnerIntervals: public NSearchMapParser::ISearchMapScannerCallback {
    private:
        TSet<NSearchMapParser::TSearchMapHost> UsefulSlotsInfo;
        TSet<NSearchMapParser::TSearchMapHost> RestoreSlotsInfo;
        TSet<TString> NotUsageSlots;
        TString ServiceName;
        TString SlotName;

    public:

        TScannerInnerIntervals(const TVector<TString>& notUsageSlots, const TString& serviceName) {
            ServiceName = serviceName;
            for (const TString& slot : notUsageSlots) {
                NRTYCluster::TSlotData sd;
                if (!NRTYCluster::TSlotData::Parse(slot, sd))
                    ythrow yexception() << "invlid not usable slot " << slot;
                NotUsageSlots.insert(sd.ShortSlotName());
            }
        }

        TVector<NSearchMapParser::TSearchMapHost> GetUsefulSlots() const {
            return { UsefulSlotsInfo.begin(), UsefulSlotsInfo.end() };
        }

        TVector<NSearchMapParser::TSearchMapHost> GetRestoreSlots() const {
            return { RestoreSlotsInfo.begin(), RestoreSlotsInfo.end() };
        }

        void AddNewSlots(const NSearchMapParser::ISlotsPool& pool) {
            const TVector<NSearchMapParser::TSearchMapHost>& slots = pool.GetSlots();
            for (ui32 i = 0; i < slots.size(); ++i)
                RestoreSlotsInfo.insert(slots[i]);
        }

        bool OnService(const NSearchMapParser::TSearchMapService& info) final {
            return ServiceName == info.Name;
        }

        void OnHost(const NSearchMapParser::TSearchMapHost& host, const NSearchMapParser::TSearchMapReplica& /*replica*/, const NSearchMapParser::TSearchMapService& /*service*/) final {
            if (!NotUsageSlots.contains(host.GetShortSlotName()) && !RestoreSlotsInfo.contains(host)) {
                DEBUG_LOG << "SLOT " << host.GetSlotName() << " detected as useful [" << host.Shards.ToString() << "]" << Endl;
                UsefulSlotsInfo.insert(host);
            }
        }
    };

    TRestoreIndexTask::TRestoreIndexTask(const NSearchMapParser::TSlotsPool& restoreSlots, const TVector<TString>& notUsageSlots,
        const TClusterTask::TCgiContext& context, NRTYDeploy::ICommonData* commonData)
        : TClusterTask(context, commonData, RESTORE_TASK_TYPE)
        , RestoreSlots(restoreSlots)
        , NotUsageSlots(notUsageSlots)
    {
    }

    bool TRestoreIndexTask::CheckTaskConstruction(NSearchMapParser::TSearchMap& ism, NSaas::TShardsDispatcher::TContext& sharding) {
        NRTYDeployInfo::IDeployComponentInfo::TPtr info = NRTYDeployInfo::IDeployComponentInfo::TFactory::Construct(ServiceType);
        if (!info) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "Incorrect service type: " + ServiceType);
            return false;
        }
        info->SetInfo(CommonData, CType);
        ism = info->SearchMap(Service);

        const NSearchMapParser::TSearchMapService* service = ism.GetService(Service);
        if (!service) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "Can't read service " + Service);
            return false;
        }

        if (!RestoreSlots.Size()) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "slots list is empty");
            return false;
        }

        sharding = service->ShardsDispatcher->GetContext();
        for (auto&& i : RestoreSlots.GetMutableSlots()) {
            NSearchMapParser::TSearchMapHost host;
            if (!service->GetHostBySlot(i.GetSlotName(), host)) {
                SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "restore slot " + i.GetSlotName() + " not in service map");
                return false;
            }
            if (!host.DisableSearch || !host.DisableIndexing) {
                SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "slot for restore " + i.GetSlotName() + " not disabled");
                return false;
            }
        }
        for (auto&& i : NotUsageSlots) {
            if (!service->SlotExists(i)) {
                SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "not useful slot " + i + " not in service map");
                return false;
            }
        }
        return true;
    }

    NRTYScript::TSlotTaskContainer TRestoreIndexTask::BuildIndexesClearing(const NRTYScript::TSlotTaskContainer& startNode) {
        TVector<NRTYScript::TTaskContainer> finishContainers;
        for (auto &i : RestoreSlots.GetSlots()) {
            NDaemonController::TAction::TPtr action = new NDaemonController::TClearIndexAction();
            NRTYScript::TSlotTaskContainer current = GetScript()->AddDependentAction(i.Name, i.ControllerPort(), "", action, startNode.GetContainer(), true);
            current = current.AddNext(*GetScript(), new NDaemonController::TRestartAction());
            finishContainers.push_back(current.GetContainer());
        }

        return GetScript()->AddDependentAction(CommonData->GetDeployManagerBalanserHost(), CommonData->GetDeployManagerBalanserPort(), CommonData->GetDeployManagerBalanserUriPrefix(), new NRTYDeploy::TLogicAction(), finishContainers);
    }

    void TRestoreIndexTask::DoBuildTask() {
        try {
            NSearchMapParser::TSearchMap ism;
            NSaas::TShardsDispatcher::TContext sharding(NSaas::KeyPrefix);
            if (!CheckTaskConstruction(ism, sharding)) {
                return;
            }

            NRTYScript::TRTYClusterScript::TPtr& script = GetScript();

            TScannerInnerIntervals scanIn(NotUsageSlots, Service);
            scanIn.AddNewSlots(RestoreSlots);

            ism.Scan(scanIn);
            NRTYReplication::TDonorsBuilder donors(scanIn.GetRestoreSlots(), scanIn.GetUsefulSlots());
            if (!donors.BuildRemap()) {
                SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, "can't build coverage for restore!!");
                return;
            }

            NRTYScript::TSlotTaskContainer start = script->AddAction(CommonData->GetDeployManagerBalanserHost(), CommonData->GetDeployManagerBalanserPort(), CommonData->GetDeployManagerBalanserUriPrefix(), new NRTYDeploy::TLogicAction());

            NRTYScript::TSlotTaskContainer buildIndexesClearingNode = BuildIndexesClearing(start);

            THashMap<TString, NRTYScript::TTaskContainer> slotsDetachingContainers;

            for (auto &i : donors.GetDonors()) {
                NDaemonController::TAction::TPtr action = new NDaemonController::TDetachAction(donors.GetDonorInfo(i), sharding, NDaemonController::apStartAndWait);
                if (donors.GetDonorInfo(i).size()) {
                    NRTYScript::TSlotTaskContainer current = script->AddDependentAction(i.Name, i.ControllerPort(), "", action, buildIndexesClearingNode,
                        true);
                    slotsDetachingContainers[i.GetSlotName()] = current.GetContainer();
                }
            }

            THashMap<TString, NRTYScript::TTaskContainer> slotsDownloadContainersStart;
            THashMap<TString, NRTYScript::TTaskContainer> slotsDownloadContainersCurrent;

            for (auto &i : donors.GetRestore()) {
                TVector<NRTYReplication::TIntervalLink> links = donors.GetRestoreInfo(i);
                for (auto &link : links) {
                    CHECK_WITH_LOG(link.Donor);
                    auto slotTask = slotsDetachingContainers.find(link.Donor->GetSlotName());
                    CHECK_WITH_LOG(slotTask != slotsDetachingContainers.end());
                    NRTYScript::TTaskContainer detachTaskContainer = slotTask->second;

                    TString idRes = "$" + detachTaskContainer.GetName() + "-action.task.id_res_" + link.Interval.ToString();
                    NDaemonController::TAction::TPtr action = new NDaemonController::TDownloadAction(idRes, NDaemonController::apStartAndWait);

                    NRTYScript::TTaskContainer current;
                    if (slotsDownloadContainersStart.contains(i.GetSlotName())) {
                        NRTYScript::TTaskContainer start = slotsDownloadContainersStart[i.GetSlotName()];
                        script->AddSeqInfo(detachTaskContainer, start, new NRTYScript::TSucceededChecker());

                        current = script->AddDependentAction(i.Name, i.ControllerPort(), "", action, slotsDownloadContainersCurrent[i.GetSlotName()], true);
                    } else {
                        current = script->AddDependentAction(i.Name, i.ControllerPort(), "", action, detachTaskContainer, true);
                        slotsDownloadContainersStart[i.GetSlotName()] = current;
                    }
                    slotsDownloadContainersCurrent[i.GetSlotName()] = current;
                }
            }

            TVector<NRTYScript::TTaskContainer> tasksDownload;
            tasksDownload.reserve(slotsDownloadContainersCurrent.size());
            for (auto&& i : slotsDownloadContainersCurrent) {
                tasksDownload.push_back(i.second);
            }
            NRTYScript::TSlotTaskContainer finish = script->AddDependentAction(CommonData->GetDeployManagerBalanserHost(), CommonData->GetDeployManagerBalanserPort(), CommonData->GetDeployManagerBalanserUriPrefix(), new NRTYDeploy::TLogicAction(), tasksDownload);
            Y_UNUSED(finish);

        } catch (...) {
            SetStatus(ctsFailed, TClusterTask::rsFailedOnConstruction, CurrentExceptionMessage());
        }
    }

};
