#include <saas/rtyserver/common/search_area_modifier.h>
#include <search/meta/context.h>
#include <search/meta/mergedres.h>
#include <search/meta/rearrange/rearrange.h>
#include <search/web/core/rule.h>

class TCacheSupporterRule : public IRearrangeRule {
    class TCacheSupporterContext : public IRearrangeRuleContext {
    public:
        inline TCacheSupporterContext(const TCacheSupporterRule* parent)
            : Parent(parent)
        {}

        void DoAdjustClientParams(const TAdjustParams& /*params*/) override {
            Disabled = LocalScheme()["Disabled"].IsTrue();
        }

        void DoRearrangeAfterMerge(TRearrangeParams& rearrangeParams) override {
            if (Disabled)
                return;

            TMetaSearchContext* ctx = rearrangeParams.Result->Context();
            TMetaGrouping* g = rearrangeParams.Current.Grouping;

            Y_ASSERT(g);

            ui64 removedDocs = 0;
            ui64 removedGroups = 0;

            for (size_t ind = 0; ind < g->Size();) {
                TMetaGroup::TDocs& docs = g->GetMetaGroup(ind).MetaDocs;
                for (size_t docInd = 0; docInd < docs.size();) {
                    const auto searchSource = ctx->GetClient(docs[docInd].Doc())->SearchSource();
                    if (!searchSource->Config()->Remote.NoCache && Parent->CheckDoc(searchSource->Num, docs[docInd]->Handle)) {
                        docs.erase(docs.begin() + docInd);
                        ++removedDocs;
                        continue;
                    }
                    docInd++;
                }

                if (!docs.size()) {
                    g->Erase(ind);
                    ++removedGroups;
                } else {
                    ind++;
                }
            }

            rearrangeParams.InsertWorkedRule("RemovedDocs", ToString(removedDocs));
            rearrangeParams.InsertWorkedRule("RemovedGroups", ToString(removedGroups));
        }

    private:
        const TCacheSupporterRule* Parent;
        bool Disabled;
    };
private:
    bool CheckDoc(const size_t& ClientNum, const TDocHandle& docId) const {
        return Singleton<TSearchAreaModifier>()->GetModifier()->IsRemoved(ClientNum, docId);
    }
public:
    inline TCacheSupporterRule(const TString& /*config*/, const TSearchConfig& /*searchConfig*/)
    {
    }

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

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

REGISTER_REARRANGE_RULE(CacheSupporter, CreateCacheSupporterRule);
