#include <saas/rtyserver/common/search_area_modifier.h>
#include <search/grouping/foundstat.h>
#include <search/meta/context.h>
#include <search/web/core/rule.h>
#include <library/cpp/logger/global/global.h>

//
//The rule allows to artificially remove a percentage of documents with greater docids from processing
//
class TRankThresholdRule final : public IRearrangeRule {
    class TRankThresholdContext final : public IRearrangeRuleContext {
    public:
        inline TRankThresholdContext(ISearchAreaModifier* sam)
            : SAM(sam)
            , SimDocCount(false)
            , Fraction(DefaultFraction)
        {
            CHECK_WITH_LOG(SAM);
        }


        void DoAdjustClientParams(const TAdjustParams& /*params*/) override {
            SimDocCount = LocalScheme()["SimDocCount"].IsTrue();
            Fraction = LocalScheme()["Fraction"].GetNumber(DefaultFraction); // CGI: &rearr=scheme_Local/RankThreshold/Fraction=0.3
        }

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

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

            Y_ASSERT(g && ctx);

            ui64 filteredDocs = 0;
            ui64 removedDocs = 0;
            ui64 removedGroups = 0;

            THolder<ISearchLocker::ILock> lock(SAM->GetSearchLock());

            for (size_t i = 0; i < g->Size();) {
                TMetaGroup::TDocs& docs = g->GetMetaGroup(i).MetaDocs;
                for (auto iDoc = docs.begin(); iDoc != docs.end();) {
                    const TMergedDoc& doc = *iDoc;

                    bool keepIt = CheckRule(ctx, doc);
                    if (keepIt) {
                        ++iDoc;
                        ++filteredDocs;
                    } else {
                        iDoc = docs.erase(iDoc);
                        ++removedDocs;
                    }
                }

                if (docs.empty()) {
                    g->Erase(i);
                    ++removedGroups;
                } else {
                    ++i;
                }
            }

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

            if (SimDocCount) {
                //Simulate the change of stats
                TFoundStatistics64 numDocs = g->NumDocs;
                TFoundStatistics numGroups = g->RelevStat;
                constexpr int numPriorities = sizeof(numDocs.Stat) / sizeof(numDocs.Stat[0]);
                for (int i = 0; i < numPriorities; ++i) {
                    numDocs.Stat[i] = (ui64) ceil(numDocs.Stat[i] * Fraction);
                    numGroups.Stat[i] = (ui32) ceil(numGroups.Stat[i] * Fraction);
                }
                g->SetRelevStats(numDocs, numGroups);
            }
        }

        inline bool CheckRule(TMetaSearchContext *ctx, const TMergedDoc &doc) {
            size_t clientNum = GetClientNum(ctx, doc);
            if (clientNum == INVALID_CLIENT_NUM)
                return false;

            ui32 docId = DecodeDocId(doc, clientNum);
            ui32 srcDocCount = SAM->GetDocumentsCount(clientNum, /*withDeleted=*/true);
            ui32 maxDocId = (srcDocCount != (ui32)-1)
                            ? (ui32) ceil(srcDocCount * Fraction)  //controller is alive
                            : std::numeric_limits<ui32>::max(); // controller is removed

            return docId <= maxDocId;
        }

        inline bool IsDisabled() const {
            return Fraction < 0;
        };

    private:
        inline static size_t GetClientNum(TMetaSearchContext* ctx, const TMergedDoc& doc) {
            const auto client = ctx->GetClient(doc.Doc());
            if (!client)
                return INVALID_CLIENT_NUM;

            const auto searchSource = client->SearchSource();
            if (!searchSource)
                return INVALID_CLIENT_NUM;

            return searchSource->Num;
        }

        inline ui32 DecodeDocId(const TMergedDoc& doc, size_t clientNum) const {
            Y_ASSERT(clientNum != INVALID_CLIENT_NUM);
            Y_ASSERT(doc->Handle.DocRoute.Length() == 0);

            return SAM->DecodeDocId(doc->Handle, clientNum); //throws an exception if failed
        }

    private:
        ISearchAreaModifier* SAM;
        bool SimDocCount;
        double Fraction;
        static constexpr double DefaultFraction = -1.f;   //disabled by default
    };


public:
    inline TRankThresholdRule(const TString& /*config*/, const TSearchConfig& /*searchConfig*/)
    {
    }

    IRearrangeRuleContext* DoConstructContext() const override {
        return new TRankThresholdContext(Singleton<TSearchAreaModifier>()->GetModifier());
    }
};

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


REGISTER_REARRANGE_RULE(RankThreshold, CreateRankThresholdRule);
