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

class TCountSortRule : public IRearrangeRule {
    class TCountSortContext : public IRearrangeRuleContext {
    public:
        void DoRearrangeAfterMerge(TRearrangeParams& rearrangeParams) override {
            TMetaGrouping* g = rearrangeParams.Current.Grouping;
            if (!strcmp(g->GetGroupMethod(), "count")) {
                TMultiMap<ui32, ui32> groups;
                for (size_t ind = 0; ind < g->Size(); ind++) {
                    groups.insert(std::make_pair(g->GetMetaGroup(ind).NumDocs[2], ind));
                }
                TVector<size_t> groupsOrdered;
                groupsOrdered.resize(g->Size(), (size_t)-1);
                if (g->SP.QAscendOrder) {
                    ui32 id = 0;
                    for (TMultiMap<ui32, ui32>::const_iterator i = groups.begin(), e = groups.end(); i != e; ++i, ++id) {
                        groupsOrdered[id] = i->second;
                    }
                }
                else {
                    ui32 id = 0;
                    for (TMultiMap<ui32, ui32>::const_reverse_iterator i = groups.rbegin(), e = groups.rend(); i != e; ++i, ++id) {
                        groupsOrdered[id] = i->second;
                    }
                }
                g->RearrangeGroups(groupsOrdered);
            }
        }
    };
public:
    inline TCountSortRule(const TString& /*config*/, const TSearchConfig& /*searchConfig*/)
    {
    }

    virtual ~TCountSortRule() {
    }

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

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

REGISTER_REARRANGE_RULE(CountSort, CreateCountSortRule);

