#pragma once

#include "bundle_model.h"
#include "timer_wrapper.h"
#include <robot/jupiter/library/rtdoc/merger/counts_merger.h>
#include <robot/jupiter/library/rtdoc/file/local_client.h>
#include <robot/jupiter/library/rtdoc/file/yt_client.h>

#include <library/cpp/logger/log.h>
#include <util/generic/hash.h>

namespace NFusion {
    class TPortionsBundle {
    public:
        using TPrepTypesTable = THashMap<TString, const NProtoBuf::Descriptor*>;

        enum class EStage {
            Undefined,
            MergePortions,
            Final
        };

    private:
        NRtDoc::TBuilderLocalClientPtr PrepClient;
        TPrepTypesTable PrepTypes;
        TLog& Log;
        TTimerLog& TimerLog;

    public:
        TPortionsBundle(NRtDoc::TBuilderLocalClientPtr client, TLog& log, TTimerLog& timerLog);

        void Run(const TMercuryCmds& allCommands);

        void RunTwice(const TMercuryPrepCmd& twoStageCommand);

        void MergePortions(const TMercuryPrepCmd& cmd, NRtDoc::TSegmentHitRemapper::TPtr remapper, const TString& publishDir, bool keepDelta);

        void ReduceMergedPortions(const TMercuryPrepCmd& cmd, const TString& publishDir);

        void RewritePortion(const TMercuryPrepCmd& cmd);

    public:
        static void RegisterPortionIds(NRtDoc::TBuilderPortionsConfig& client);

        static void RegisterDocIdSetters(NRtDoc::ILocalClient& client);

        static void AddDocIdSettersForStage(NRtDoc::ILocalClient& client, EStage portionsStage);
    };
}
