#pragma once

#include "text_machine_plugin.h"

#include <saas/rtyserver/factors/feature.h>

#include <search/lingboost/relev/text_machine_features.h>

#include <kernel/factor_storage/factor_view.h>

#include <kernel/reqbundle_iterator/reqbundle_iterator_yndex_builder.h>


class TSaasTextMachineFeaturesCalcerFactory : public TTextMachineFeaturesCalcerFactoryIteratorsBase {
    TIntrusivePtr<ISaasTextMachinePlugin> Engine;
    size_t ReqBundleNumElems;
    TVector<ui64> Kps;
public:
    TSaasTextMachineFeaturesCalcerFactory(
        TIntrusivePtr<ISaasTextMachinePlugin> engine,
        NReqBundle::TConstReqBundleAcc reqBundle,
        const NLingBoost::TRawConfig& config,
        const TReqBundleIteratorsHashers& hashers,
        TVector<ui64> kps)
        : TTextMachineFeaturesCalcerFactoryIteratorsBase(reqBundle, hashers, config)
        , Engine(engine)
        , ReqBundleNumElems(reqBundle.GetSequence().GetNumElems())
        , Kps(std::move(kps))
    {};

    THolder<NLingBoost::IProductionFeaturesCalcer> CreateProductionCalcer() const override {
        return Engine->CreateProductionCalcer(GetTargetTextMachineOptions());
    }

    THolder<NLingBoost::IExperimentFeaturesCalcer> CreateExperimentCalcer() const override {
        return Engine->CreateExperimentCalcer(GetTargetTextMachineOptions()); // may return nullptr
    }

    void OpenTrIterator(const TYndexRequester* textYR,
        TLangMask langMask,TLangMask flatBastardsLangMask,
        const ISentenceLengthsLenReader* txtSentReader, const ISentenceZonesReader* sentenceZonesReader,
        const ui32 maxTitleBreak, const TSentenceZones titleZones) {

        ConfigureBlockTypes(NLingBoost::TBaseIndex::Text, GetTargetData().GoodForIndex[NLingBoost::TBaseIndex::Text]);
        if (textYR) {
            TReqBundleYndexIteratorBuilder builder(*textYR, IHitsAccess::IDX_TEXT, langMask, flatBastardsLangMask);
            builder.SetKeyPrefixes(Kps);
            GetTargetData().ReqBundleTrHitsProvider = Base.OpenIterator(builder, TrHasher);
            if (txtSentReader) {
                GetTargetData().TextHitsStreamizer = NLingBoost::CreateSaasTextHitsStreamizer(txtSentReader, sentenceZonesReader, maxTitleBreak, titleZones);
                GetTargetData().TextHitsStreamizer->TolerateErrors(NLingBoost::EStreamizerErrorType::BreakProcessingError);
                GetTargetData().TextHitsStreamizer->TolerateErrors(NLingBoost::EStreamizerErrorType::ZeroBasedHitFromIterator);
            }
        }
    }

    void InitTrHitsProviderData(const ISentenceLengthsLenReader* txtSentReader, const ISentenceZonesReader* sentenceZonesReader,
        const ui32 maxTitleBreak, const TSentenceZones titleZones) {
        if (!txtSentReader)
            return;
        GetTargetData().HasHitsProviderExternalData[NLingBoost::TBaseIndex::Text] = true;
        ConfigureBlockTypes(NLingBoost::TBaseIndex::Text, GetTargetData().GoodForIndex[NLingBoost::TBaseIndex::Text]);
        GetTargetData().TextHitsStreamizer = NLingBoost::CreateSaasTextHitsStreamizer(txtSentReader, sentenceZonesReader, maxTitleBreak, titleZones);
        GetTargetData().TextHitsStreamizer->TolerateErrors(NLingBoost::EStreamizerErrorType::BreakProcessingError);
        GetTargetData().TextHitsStreamizer->TolerateErrors(NLingBoost::EStreamizerErrorType::ZeroBasedHitFromIterator);
    }

    void OpenAnnIterator(const TYndexRequester* annYR,
        TLangMask langMask, TLangMask flatBastardsLangMask, TCateg relevCountry,
        const NIndexAnn::IDocDataIndex* annIndex, const TSentenceLengthsReader* annSentReader)
    {
        Y_VERIFY(annYR);
        ConfigureBlockTypes(NLingBoost::TBaseIndex::Ann, GetTargetData().GoodForIndex[NLingBoost::TBaseIndex::Ann]);
        TReqBundleYndexIteratorBuilder builder(*annYR, IHitsAccess::IDX_TEXT, langMask, flatBastardsLangMask, annSentReader);
        builder.SetKeyPrefixes(Kps);
        GetTargetData().ReqBundleAnnHitsProvider = Base.OpenIterator(builder, AnnHasher);
        GetTargetData().AnnHitsStreamizer = Engine->CreateAnnHitsStreamizer(relevCountry, annIndex, annSentReader);
    }

    void InitAnnHitsProviderData(TCateg relevCountry, const NIndexAnn::IDocDataIndex* annIndex, const TSentenceLengthsReader* annSentReader) {
        GetTargetData().HasHitsProviderExternalData[NLingBoost::TBaseIndex::Ann] = true;
        ConfigureBlockTypes(NLingBoost::TBaseIndex::Ann, GetTargetData().GoodForIndex[NLingBoost::TBaseIndex::Ann]);
        GetTargetData().AnnHitsStreamizer = Engine->CreateAnnHitsStreamizer(relevCountry, annIndex, annSentReader);
    }

    void OpenFactorAnnIterator(const TYndexRequester* factorAnnYR, TCateg relevCountry, const TSentenceLengthsReader* factorAnnSentLenReader,
        const NIndexAnn::IDocDataIndex* factorAnnIndex, ui32 factorannFirstStageFetchLimit)
    {
        Y_VERIFY(factorAnnYR);
        ConfigureBlockTypes(NLingBoost::TBaseIndex::FactorAnn, GetTargetData().GoodForIndex[NLingBoost::TBaseIndex::FactorAnn]);
        TReqBundleIteratorOptions options(NLingBoost::TBlockType::GetAllLayersMask());
        options.IteratorsFirstStageFetchLimit = factorannFirstStageFetchLimit;

        TReqBundleYndexIteratorBuilder builder(*factorAnnYR, IHitsAccess::IDX_TEXT, THitsLoader::GetDefaultLemmatizedLangMask(), THitsLoader::GetDefaultFlatBastardsLangMask(), factorAnnSentLenReader);
        builder.SetKeyPrefixes(Kps);
        GetTargetData().ReqBundleFactorAnnHitsProvider = Base.OpenIterator(builder, FactorAnnHasher, options);
        GetTargetData().FactorAnnHitsStreamizer = Engine->CreateFactorAnnHitsStreamizer(relevCountry, factorAnnIndex, factorAnnSentLenReader);
    }
};

class TSaasTextMachineFeaturesCalcer final : public NRTYFeatures::ISlicedCalcer {
    THolder<TTextMachineFeaturesCalcer> Calcer;
    NFactorSlices::EFactorSlice SliceId;
    TFactorView ViewMain;
    TFactorView ViewOnotole;
    ui32 DbgLastDocId;

public:
    TSaasTextMachineFeaturesCalcer(THolder<TTextMachineFeaturesCalcer> calcer, NFactorSlices::EFactorSlice sliceId)
        : Calcer(std::move(calcer))
        , SliceId(sliceId)
        , ViewMain(0)
        , ViewOnotole(0)
        , DbgLastDocId(Max<ui32>())
    {
        Y_VERIFY(Calcer);
    }

    void Bind(TFactorStorage& storage) {
        Y_VERIFY(!ViewMain && !ViewOnotole, "new calcer must be created before each pass");

        ViewMain = storage.CreateViewFor(SliceId);
        ViewOnotole = storage.CreateViewFor(EFactorSlice::LINGBOOST_ONOTOLE);

        if (ViewMain.Size() == 0) {
            ythrow yexception() << "Slice should be enabled: " << SliceId;
        }
    }

    void CalcFeatures(ui32 docId) {
        // Note(yrum): the assertion below is to avoid nasty bugs like SAASSUP-2543 (underlying KeyInv iterators may return no hits when docId sequence
        // is not monotonic). The text machine itself has no such restriction, and if it were run over the modernized keyinv.wad, we could have this restriction removed.
        Y_ASSERT(DbgLastDocId == Max<ui32>() || DbgLastDocId < docId);
        Y_IF_DEBUG(DbgLastDocId = docId);

        Calcer->CalcFeatures(docId, ViewMain, ViewOnotole);
    }

    const TTextMachineFeaturesCalcer& GetStats() const {
        // this accessor is intended for info requests only, Calc() methods are non-const
        return *Calcer;
    }
};

struct TSaasTextMachineCalcerHolder {
    THolder<TTextMachineFeaturesCalcer> WebProductionCalcer;
    TIntrusivePtr<TSaasTextMachineFeaturesCalcer> CustomCalcer;

public:
    TSaasTextMachineCalcerHolder() = default;

    TSaasTextMachineCalcerHolder(THolder<TTextMachineFeaturesCalcer> webProdCalcer)
        : WebProductionCalcer(std::move(webProdCalcer))
    {
    }

    TSaasTextMachineCalcerHolder(TIntrusivePtr<TSaasTextMachineFeaturesCalcer> customCalcer)
        : CustomCalcer(std::move(customCalcer))
    {
    }

    bool IsCustom() const {
        Y_ASSERT(!WebProductionCalcer || !CustomCalcer);
        return !!CustomCalcer;
    }
};
