#include <saas/rtyserver/common/sharding.h>

#include <search/meta/context.h>
#include <search/meta/mergedres.h>
#include <search/meta/rearrange/rearrange.h>
#include <search/web/core/rule.h>

#include <kernel/qtree/richrequest/richnode.h>

class TKeyPrefixSupporterRule : public IRearrangeRule {
public:
    class TKeyPrefixSupporterRuleContext : public IRearrangeRuleContext {
    public:
        TKeyPrefixSupporterRuleContext(const TKeyPrefixSupporterRule& parent)
            : Parent(parent)
        {}
        virtual ~TKeyPrefixSupporterRuleContext() {}

        void DoPrepareRearrangeParams(const TAdjustRuleParams& arp) override {
            PreparePatchedQtrees(arp);
        }

        void DoPreparePassageRearrangeParams(const TAdjustRuleParams& arp) override {
            PreparePatchedQtrees(arp);
        }

        void DoAdjustClientParams(const TAdjustParams& ap) override {
            AdjustRequest(ap);
        }

        void DoAdjustFactorRequests(const TAdjustParams& ap) override {
            AdjustRequest(ap);
        }

    private:
        void PreparePatchedQtrees(const TAdjustRuleParams& arp) {
            const i64 shardsNumber = LocalScheme()["ShardsNumber"].GetIntNumber();

            TVector<TString> kps;
            Split(arp.RequestFields.CgiParam.Get("kps"), ",", kps);
            if (arp.RP.HasRequestTree()) {
                TRichTreePtr originalRichTree = arp.RP.GetRequestTree()->Copy();
                for (int kp = 0; kp < kps.ysize(); kp++) {
                    if (kps[kp] != "0")
                        KpsShards[NRTYServer::GetShard(FromString(kps[kp]), static_cast<int>(shardsNumber))].push_back(FromString(kps[kp]));
                }

                for (TMap<ui32, TVector<ui64> >::const_iterator i = KpsShards.begin(), e = KpsShards.end(); i != e; ++i) {
                    originalRichTree->Root->UpdateKeyPrefix(i->second);
                    TBinaryRichTree buf;
                    originalRichTree.Get()->Serialize(buf);
                    TString NewQTree = EncodeRichTreeBase64(buf);
                    KpsTasks[i->first].push_back(NewQTree);
                }
            }
        }

        void ReplaceParam(TMetaRequestAdjusterRef& client, const char* name, const TString& newValue) {
            for (int i = client->ClientFormFieldCount(name) - 1; i >= 0; i--)
                client->ClientFormFieldRemove(name, i);

            client->ClientFormFieldInsert(name, newValue.c_str());
        }

        void RemoveParam(TMetaRequestAdjusterRef& client, const char* name) {
            for (int i = client->ClientFormFieldCount(name) - 1; i >= 0; i--)
                client->ClientFormFieldRemove(name, i);
        }

        void AdjustRequest(const TAdjustParams& ap) {
            if (ap.ClientRequestAdjuster->IsClientEphemeral())
                return;

            if (!KpsTasks.size())
                return;

            int numQuery = 0;
            const NRTYServer::TShard shardCurrent = NRTYServer::GetShard(ap.ClientRequestAdjuster->SourceGroup());
            const TVector<TString>& kps = KpsTasks[shardCurrent];
            for (int kp = 0; kp < kps.ysize(); kp++) {
                TMetaRequestAdjusterRef clientRequestAjuster = nullptr;
                if (!numQuery)
                    clientRequestAjuster = ap.ClientRequestAdjuster;
                else
                    clientRequestAjuster = ap.ClientRequestAdjuster->ClientAddRequest();

                if (!!clientRequestAjuster) {
                    RemoveParam(clientRequestAjuster, "kps");
                    RemoveParam(clientRequestAjuster, "mss");
                    ReplaceParam(clientRequestAjuster, "qtree", kps[kp]);
                    numQuery++;
                }
            }

            if (!numQuery)
                ap.ClientRequestAdjuster->ClientDontSendRequest();
        }

    protected:
        const TKeyPrefixSupporterRule& Parent;
        TMap<size_t, TVector<TString>> KpsTasks;
        TMap<ui32, TVector<ui64> > KpsShards;
    };
public:
    TKeyPrefixSupporterRule(const TString& /*config*/, const TSearchConfig& searchConfig)
        : SearchConfig(searchConfig)
    {}
    virtual ~TKeyPrefixSupporterRule() {}

    IRearrangeRuleContext* DoConstructContext() const override {
        return new TKeyPrefixSupporterRuleContext(*this);
    }

    const TSearchConfig& GetSearchConfig() const {
        return SearchConfig;
    }
protected:
    const TSearchConfig& SearchConfig;
};

IRearrangeRule* CreateKeyPrefixSupporterRule(const TString& config, const TSearchConfig& searchConfig) {
    return new TKeyPrefixSupporterRule(config, searchConfig);
}

REGISTER_REARRANGE_RULE(KeyPrefixSupporter, CreateKeyPrefixSupporterRule);
