#pragma once

#include "fresh_pruning_builder.h"

#include <robot/library/oxygen/indexer/processor/pruning/pruning.h>
#include <robot/library/oxygen/indexer/library/pruning_builder/pruning_builder.h>
#include <robot/library/oxygen/indexer/processor/processor/processor.h>
#include <robot/library/oxygen/indexer/processor/protos/config.pb.h>
#include <yweb/robot/urlzones/filters.h>

#include <util/thread/pool.h>
#include <util/generic/maybe.h>

//
// This processor calculates the final DocId values of all documents
// and dumps groups and pruning data files.
//
// See: http://wiki.yandex-team.ru/JandeksPoisk/Oxygen/Tech/Prunning.
//
class TPruningDataAccessor;
struct TPruningDataSelectionRank;
struct TPruningZoneParams;

namespace NZones {
    class TLangDomainSearchZoneFilter;
}

class TFreshMixer : public NOxygen::TPruningProcessorBase {
private:
    struct TZonePruningOpts {
        ELanguage LanguageRegion;
        const TString RankValuesTupleName;
        const TString SecondaryRankValuesTupleName;
        const TString SelectionRank; // Selection rank is the rank to use from rank values
        const TString SecondarySelectionRank; // Selection rank is the rank to use from rank values
        THolder<NOxygen::TPruningProcessor::ISelectionRankSource> SelectionRankSource;
        THolder<NOxygen::TPruningProcessor::ISelectionRankSource> SecondarySelectionRankSource;

        ui32 Proportion;
        bool IsRequired;
    };

    const TAtomicSharedPtr<IThreadPool> TaskPool;

    const TString TupleNameForUrl;
    const TString TupleNameForMixerLanguage; //Tuple name for the document language value
    const TString OutputPrefix;
    const NOxygen::TPruningConfig PruningConfig;

    THolder<TPruningDataAccessor> PruningData;
    THolder<NZones::TLangDomainSearchZoneFilter> LangDomainFilter;
    THolder<NZones::TLangSearchZoneFilter> LangSearchFilter;
    TMaybe<NOxygen::TDocIdMap> FinalMap;
    NOxygen::TFuture<void> AsyncDelay;

    static NOxygen::TPruningConfig LoadPruningConfig(const NOxygen::TRTYMixerOptions &options);
    TPruningZoneParams BuildZoneParams() const;

    TMaybe<float> TryFetchRank(TStringBuf url, NOxygen::TObjectContext& objectContext, const TZonePruningOpts& zone) const;
    bool FetchRanks(TStringBuf url, TRankValues& ranks, NOxygen::TObjectContext& objectContext);
    bool BelongsToLangZone(ELanguage languageZone, const TString& host, ELanguage docLanguage) const;

    TVector<TZonePruningOpts> ZoneOpts;

protected:
    THolder<NOxygen::IPruningBuilder> CreatePruningBuilder() const override;

public:
    TFreshMixer(
            const TAtomicSharedPtr<IThreadPool> taskPool,
            const NOxygen::TProcessorPtr slave,
            const NOxygen::TRTYMixerOptions &options,
            const TString& outputPrefix,
            bool selectionRankStore = false);
    ~TFreshMixer();

    TString GetClassName() const override {
        return "TFreshMixer";
    }

    // ITupleProcessor.
    NOxygen::TTuplesUsageInfo GetRequiredTuples() const override;
    void Start() override;
    NThreading::TFuture<NOxygen::TReturnObjectContext> Process(NOxygen::TObjectContext objectContext, ui32 tmpObjectId) override;
    void Finish(const NOxygen::TDocIdMap* externalMap) override;

    [[nodiscard]] NOxygen::TDocIdMap CreateDocIdMapAsInFinish(const NOxygen::TDocIdMap* externalMap) const override;

    // shall be called after Finish
    const NOxygen::TDocIdMap& GetDocIdMap() const override;

    static NOxygen::TPruningProcessorBase::ESourceType ProtoToSourceType(NOxygen::TRTYMixerOptions::ESourceType);
};
