#include <saas/searchproxy/search_meta/context.h>
#include <search/web/core/rule.h>
#include <search/meta/context.h>
#include <saas/searchproxy/common/cgi.h>
#include <saas/rtyserver/factors/rank_model.h>
#include "l3.h"

namespace NSaas {
    TString EncodeFactorsIndices(const NRTYFactors::TUsedFactors& factors) {
        TStringStream result;
        constexpr ui32 INVALID_FACTOR_IDX = Max<ui32>();
        ui32 prev = INVALID_FACTOR_IDX;
        ui32 intervalStart = INVALID_FACTOR_IDX;
        for (const ui32 factorIdx: factors) {
            if (prev == INVALID_FACTOR_IDX || factorIdx != prev + 1) {
                if (prev != INVALID_FACTOR_IDX) {
                    if (prev != intervalStart) {
                        result << '-' << prev;
                    }
                    result << ',';
                }
                result << factorIdx;
                intervalStart = factorIdx;
            }
            prev = factorIdx;
        }
        if (prev != intervalStart) {
            result << '-';
            result << prev;
        }
        return result.Str();
    }
}

class TRTYL3RuleContext: public IRearrangeRuleContext {
private:
    bool IsDisabled() {
        return LocalScheme().Get("Disable").IsTrue();
    }

    void AddFactorIdxsParam(const TAdjustParams& ap, bool isFactorRequest) {
        if (IsDisabled())
            return;

        const auto* ctx = dynamic_cast<const TCgiAwareSearchContext*>(&ap.ReqResults);
        if (!ctx || !ctx->MetaRankingOptions().Enabled() || ctx->MetaRankingOptions().NoFactorsRequest == isFactorRequest) {
            return;
        }

        TAtomicSharedPtr<NRTYFactors::TConfig> factors = ctx->GetServiceConfig().GetFactorsConfig();
        if (!factors) {
            return;
        }

        const NRTYFactors::TRankModelHolder* rankModel = factors->GetL3RankModel(NSearchProxyCgi::GetL3FormulaName(ctx->CgiParam));
        if (rankModel == nullptr || rankModel->GetUsedFactors().empty()) {
            return;
        }

        TString factorIndices = NSaas::EncodeFactorsIndices(rankModel->GetUsedFactors());
        ap.ClientRequestAdjuster->ClientAppendPron("dont_check_restricts", "da");
        ap.ClientRequestAdjuster->ClientAppendRelev("factors_indices", factorIndices);
    }

    size_t Iteration = 0;

public:
    TRTYL3RuleContext() {}

    // before search
    void DoAdjustClientParams(const TAdjustParams& ap) override {
        AddFactorIdxsParam(ap, false);
    }

    // after search, before ranking
    void DoAdjustFactorRequests(const TAdjustParams& ap) override {
        AddFactorIdxsParam(ap, true);
    }

    // after ranking and merge, before snippets
    void DoRearrangeAfterMerge(TRearrangeParams& rp) override {
        if (IsDisabled())
            return;

        Iteration += 1;
        const size_t maxCycles = LocalScheme().Get("MaxCycles").GetIntNumber(100500);
        if (Iteration >= maxCycles) {
            for (const auto& grouping : rp.Result->GroupingMap) {
                grouping.second->FixCount();
            }
        }
    }
};

struct TRTYL3Rule: IRearrangeRule {

    TRTYL3Rule(TString options) {
        Y_UNUSED(options);
    }

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

IRearrangeRule* CreateRTYL3Rule(const TString& options, const TSearchConfig&) {
    const NRearrConf::TRearrangeRuleParams params(options);

    return new TRTYL3Rule(options);
}

REGISTER_REARRANGE_RULE(RTYL3, CreateRTYL3Rule);
