#include "dc_alias_filter.h"

#include <saas/deploy_manager/meta/cluster.h>
#include <saas/util/external/dc.h>

namespace NRTYDeploy {
    bool TDCAliasFilter::ModifyHost(NSearchMapParser::TSearchMapHost& host, const NSearchMapParser::TSearchMapReplica& originalReplica, const NSearchMapParser::TSearchMapService& originalService) {
        NRTYCluster::TSlotData slotData(host.Name, host.SearchPort);
        TString alias;
        for (int i = 0; (!alias || alias == ALIAS_NOT_DEFINED) && i < 5; ++i) {
            alias = slotData.GetDCAlias(true);
        }
        if (!alias) {
            alias = ALIAS_NOT_DEFINED;
        }
        if (originalService.WeightsByAliasFrom.empty())
            return ModifyHostAutoWeights(host, originalReplica, originalService, alias);
        else
            return ModifyHostCustomWeights(host, originalReplica, originalService, alias);
    }

    bool TDCAliasFilter::ModifyEndpointSet(NSearchMapParser::TSearchMapHost& host, const NSearchMapParser::TSearchMapReplica& originalReplica, const NSearchMapParser::TSearchMapService& originalService) {
        TDatacenterUtil::Instance().SetDatacenter(host.Name, host.DC);
        return ModifyHost(host, originalReplica, originalService);
    }

    bool TDCAliasFilter::ModifyHostAutoWeights(NSearchMapParser::TSearchMapHost& host, const NSearchMapParser::TSearchMapReplica& originalReplica, const NSearchMapParser::TSearchMapService& originalService, const TString& alias) {
        TSet<NSearchMapParser::TSearchMapHost> hostsSortMainAlias;
        TSet<NSearchMapParser::TSearchMapHost> hostsSortOtherAlias;
        NRTYCluster::TSlotData slotData(host.Name, host.SearchPort);
        for (auto&& i : originalReplica.GetSlots()) {
            if (i.Shards == host.Shards) {
                NRTYCluster::TSlotData currentSlotData(i.Name, i.SearchPort);
                if (currentSlotData.GetDCAlias(true) == Alias) {
                    hostsSortMainAlias.insert(i);
                } else {
                    hostsSortOtherAlias.insert(i);
                }
            }
        }

        TString weights;

        bool isMainAlias = slotData.GetDCAlias() == Alias;

        for (ui32 i = 0; i < hostsSortMainAlias.size(); ++i) {
            weights += isMainAlias ? "1" : "0";
            weights += "@";
        }

        for (ui32 i = 0; i < hostsSortOtherAlias.size(); ++i) {
            weights += isMainAlias ? "0" : "1";
            weights += "@";
        }

        weights += "1";

        if (!originalService.PerDcSearch) {
            host.Group = alias + "-" + weights;
            return true;
        } else {
            host.Group = alias + "-1";
        }
        return Alias == alias;
    }

    bool TDCAliasFilter::ModifyHostCustomWeights(NSearchMapParser::TSearchMapHost& host, const NSearchMapParser::TSearchMapReplica& originalReplica, const NSearchMapParser::TSearchMapService& originalService, const TString& alias) {
        auto iterFrom = originalService.WeightsByAliasFrom.find(Alias);
        if (iterFrom == originalService.WeightsByAliasFrom.end()) {
            return ModifyHostAutoWeights(host, originalReplica, originalService, alias);
        }
        auto iterTo = iterFrom->second.find(alias);
        if (iterTo == iterFrom->second.end()) {
            return ModifyHostAutoWeights(host, originalReplica, originalService, alias);
        }
        host.Group = alias + "-" + ToString(iterTo->second);
        return true;
    }

    bool TDCAliasFilter::ModifyReplica(NSearchMapParser::TSearchMapReplica& replica, const NSearchMapParser::TSearchMapReplica& originalReplica, const NSearchMapParser::TSearchMapService& /*originalService*/) {
        if (replica.Hosts.empty())
            replica.Hosts = originalReplica.Hosts;
        return true;
    }

    TDCAliasFilter::TDCAliasFilter(const TString& alias)
        : Alias(alias ? alias : ALIAS_NOT_DEFINED)
    {}
} //namespace NRTYDeploy
