#include "portions_bundle.h"
#include "log.h"

#include <saas/rtyserver_jupi/library/tables/tables.h>

#include <robot/jupiter/library/rtdoc/file/yt_client.h>
#include <robot/jupiter/library/tables/shards_prepare.h>

#include <robot/jupiter/protos/calculated_factors.pb.h>
#include <robot/jupiter/protos/external_relev_attrs.pb.h>
#include <robot/jupiter/protos/hnsw.pb.h>
#include <robot/jupiter/protos/inv_url.pb.h>
#include <robot/jupiter/protos/offroad_keyinv.pb.h>
#include <robot/jupiter/protos/panther_dssm_doc.pb.h>
#include <robot/jupiter/protos/walrus.pb.h>
#include <robot/jupiter/protos/itditp_slim_index_data.pb.h>

#include <robot/jupiter/library/rtdoc/merger/counts_merger.h>

#include <robot/jupiter/tools/shards_prepare/lib/build_panther_dssm_embeddings.h>
#include <robot/jupiter/tools/shards_prepare/lib/panther_utils/doc_panther_terms_builder.h>
#include <robot/jupiter/tools/shards_prepare/lib/panther_utils/index_part.h>

#include <library/cpp/logger/global/global.h>

#define LOG C_LOG(this->Log)

namespace NFusion {
    template <typename TMessage, bool SkipEmpty>
    class TSimpleDocIdPreprocessorBase: public NRtDoc::IYtReaderRowPreprocessor {
    public:
        void ProcessRow(google::protobuf::Message* msg, ui32 docId) const final {
            Y_ASSERT(msg && msg->GetDescriptor() == GetRowType());
            TMessage& row = static_cast<TMessage&>(*msg);
            if constexpr (SkipEmpty) {
                // for some data types, .HasDocId() is used to differentiate between null and non-null records.
                // For the non-null record, the value assigned to DocId is bogus ('0' is used for all rows).
                if (Y_UNLIKELY(!row.HasDocId())) {
                    Y_VERIFY_DEBUG(row.SerializeAsString() == TMessage().SerializeAsString(),
                        "non-empty row is expected to have a DocId");
                    return;
                }
            }
            row.SetDocId(docId);
            ProcessTypedRow(row, docId);
        }

        const google::protobuf::Descriptor* GetRowType() const final {
            return TMessage::descriptor();
        }

        THolder<google::protobuf::Message> CreateEmptyObject() const final {
            return MakeHolder<TMessage>();
        }

        ui32 GetDocId(const google::protobuf::Message& msg) const final {
            Y_ASSERT(msg.GetDescriptor() == GetRowType());
            return static_cast<const TMessage&>(msg).GetDocId();
        }

    protected:
        virtual void ProcessTypedRow(TMessage&, ui32) const {
        }
    };

    template <typename TMessage>
    using TSimpleDocIdPreprocessor = TSimpleDocIdPreprocessorBase<TMessage, false>;

    template <typename TMessage>
    using TOptionalDataDocIdPreprocessor = TSimpleDocIdPreprocessorBase<TMessage, true>;

    class TWalrusDocId final: public TSimpleDocIdPreprocessor<NJupiter::TWalrus> {
    protected:
        void ProcessTypedRow(NJupiter::TWalrus& walrus, ui32 docId) const override {
            walrus.SetLocalDocId(docId);
            walrus.SetChunk(0);
        }
    };

    class TCalcAttrsDocId final: public TSimpleDocIdPreprocessor<NJupiter::TCalculatedFactors> {
    protected:
        void ProcessTypedRow(NJupiter::TCalculatedFactors& calcAttrs, ui32 docId) const override {
            calcAttrs.SetLocalDocId(docId);
            calcAttrs.SetChunk(0);
        }
    };

    using TExternalRelevAttrsDocId = TSimpleDocIdPreprocessor<NJupiter::TExternalRelevAttrs>;
    using THostAttrsDocId = TSimpleDocIdPreprocessor<NJupiter::THostAttrs>;
    using THnswDocEmbedDocId = TSimpleDocIdPreprocessor<NJupiter::THnswDocEmbed>;
    using TInvUrlHashDocId = TSimpleDocIdPreprocessor<NJupiter::TInvUrlHash>;
    using TArcDocId = TSimpleDocIdPreprocessor<NJupiter::TArc>;
    using TContentAttrsDocId = TSimpleDocIdPreprocessor<NJupiter::TContentAttrs>;
    using TUrlHeavyDataDocId = TSimpleDocIdPreprocessor<NJupiter::TUrlHeavyData>;
    using TNavSrcStrictDocId = TOptionalDataDocIdPreprocessor<NJupiter::TNavSrcStrict>;
    using TMobileUrlDocId = TOptionalDataDocIdPreprocessor<NJupiter::TMobileUrl>;

    template <bool IsOneDoc>
    class TPantherDssmDocsBatchDocId final: public NRtDoc::IYtReaderRowPreprocessor {
    private:
        using TMessage = NJupiter::TPantherDssmDocsBatch;
        const ui32 DocCount;

    public:
        TPantherDssmDocsBatchDocId(ui32 docCount)
            : DocCount(docCount)
        {
        }

        bool IsDocIdSetter() const override {
            return IsOneDoc;
        }

        void ProcessRow(google::protobuf::Message* msg, ui32 docId) const override {
            Y_ENSURE(DocCount > 0, "Document count must be > 0");
            Y_ASSERT(msg && msg->GetDescriptor() == GetRowType());
            TMessage& row = static_cast<TMessage&>(*msg);
            row.SetBatch(Max<ui32>() >> 1);
            if constexpr (IsOneDoc) {
                Y_ENSURE(row.GetDocs().DocsSize() == 1, "there should be only 1 doc");
                for (auto& doc : *(row.MutableDocs()->MutableDocs())) {
                    doc.SetDocId(docId);
                }
            }
        }

        const google::protobuf::Descriptor* GetRowType() const override {
            return TMessage::descriptor();
        }

        THolder<google::protobuf::Message> CreateEmptyObject() const override {
            return MakeHolder<TMessage>();
        }

        ui32 GetDocId(const google::protobuf::Message& msg) const override {
            auto& row = static_cast<const TMessage&>(msg);
            if constexpr (IsOneDoc) {
                Y_ENSURE(row.GetDocs().DocsSize() == 1, "there should be only 1 doc");
                return row.GetDocs().GetDocs(0).GetDocId();
            }
            return Max<ui32>();
        }
    };

    class TOffroadKeyInvEntryIndexId: public NRtDoc::IYtReaderRowPreprocessor {
    private:
        using TMessage = NJupiter::TOffroadKeyInvEntry;
    public:
        bool IsDocIdSetter() const override {
            return false;
        }

        void ProcessRow(google::protobuf::Message* msg, ui32 generatedDocId) const override {
            Y_ASSERT(msg && msg->GetDescriptor() == GetRowType());
            TMessage& row = static_cast<TMessage&>(*msg);
            // Trick Jupiter into thinking it has only 1 part of keyinvs, and only 1 range
            ui32 id;
            ui32 part;
            NJupiter::NPrivate::TIndexIdPart::Unpack(row.GetPart(), &id, &part);
            Y_ASSERT(part == 0);
            row.ClearFileType();
            row.ClearStartDocId();
            row.ClearPartIndex();
            row.SetPart(NJupiter::NPrivate::TIndexIdPart::Pack(generatedDocId, part));
            row.SetRange(0);
            row.SetJob(Max<ui32>() >> 1); // magic value to disable some asserts
        }

        const google::protobuf::Descriptor* GetRowType() const override {
            return TMessage::descriptor();
        }

        THolder<google::protobuf::Message> CreateEmptyObject() const override {
            return MakeHolder<TMessage>();
        }

        ui32 GetDocId(const google::protobuf::Message&) const override {
            return Max<ui32>();
        }
    };

    class TOffroadKeyInvOneDocId: public NRtDoc::IYtReaderRowPreprocessor {
    private:
        using TMessage = NJupiter::TOffroadKeyInvEntry;
    private:
        static constexpr auto BitsForRange = NJupiter::TPortionHitRemapper::MaxBitsPerJupiterKeyRange;
    private:
        mutable ui32 LastDocId = Max<ui32>();
        mutable ui32 LastPortionNo = 0;

    public:
        bool IsDocIdSetter() const override {
            return true;
        }

        void ProcessRow(google::protobuf::Message* msg, ui32 generatedDocId) const override {
            Y_ASSERT(msg && msg->GetDescriptor() == GetRowType());
            TMessage& row = static_cast<TMessage&>(*msg);
            ui32 id;
            ui32 part;
            NJupiter::NPrivate::TIndexIdPart::Unpack(row.GetPart(), &id, &part);
            Y_ENSURE(part == 0);
            row.ClearFileType();
            row.ClearStartDocId();
            row.ClearPartIndex();

            if (LastDocId != generatedDocId) {
                LastDocId = generatedDocId;
                LastPortionNo = 0;
            } else {
                Y_ASSERT(generatedDocId != Max<ui32>());
                LastPortionNo++;
            }
            Y_ENSURE(generatedDocId < (Max<ui32>() << BitsForRange) && LastPortionNo < (1 << BitsForRange));
            row.SetPart(NJupiter::NPrivate::TIndexIdPart::Pack((generatedDocId << BitsForRange) | LastPortionNo, part));
            row.SetRange(0);
            row.SetJob(Max<ui32>() >> 1); // magic value to disable some asserts
        }

        const google::protobuf::Descriptor* GetRowType() const override {
            return TMessage::descriptor();
        }

        THolder<google::protobuf::Message> CreateEmptyObject() const override {
            return MakeHolder<TMessage>();
        }

        ui32 GetDocId(const google::protobuf::Message&) const override {
            return Max<ui32>();
        }
    };
    class TOffroadKeyInvPortionIndexId: public NRtDoc::IYtReaderRowPreprocessor {
    private:
        using TMessage = NJupiter::TOffroadKeyInvEntry;
    private:
        const bool ForceRenumCounts;
        const bool OverrideRangeAndJob;
    private:
        mutable ui32 LastIndexId = Max<ui32>();
        mutable ui32 LastPart = Max<ui32>();
        mutable ui32 LastNewIndexId = Max<ui32>();
    public:
        TOffroadKeyInvPortionIndexId(const bool forceRenumCounts = false, const bool overrideRangeAndJob = false)
            : ForceRenumCounts(forceRenumCounts)
            , OverrideRangeAndJob(overrideRangeAndJob)
        {
        }

        bool IsDocIdSetter() const override {
            return false;
        }

        void ProcessRow(google::protobuf::Message* msg, ui32 generatedDocId) const override {
            Y_ASSERT(msg && msg->GetDescriptor() == GetRowType());
            Y_ENSURE(generatedDocId != Max<ui32>());
            TMessage& row = static_cast<TMessage&>(*msg);
            ui32 id;
            ui32 part;
            NJupiter::NPrivate::TIndexIdPart::Unpack(row.GetPart(), &id, &part);
            row.ClearFileType();
            row.ClearStartDocId();
            row.ClearPartIndex();

            if (id == Max<ui32>() && !ForceRenumCounts) {
                row.SetJob(Max<ui32>() >> 1); // magic value to disable some asserts
                return; // for counts_keyinv
            }

            const bool anew = id != LastIndexId || LastPart == Max<ui32>() || part != LastPart + 1;
            if (anew) {
                Y_ASSERT(part == 0); // assert the sequence is not broken - expect troubles otherwise
                LastIndexId = id;
                LastPart = part;
                LastNewIndexId = generatedDocId;
                id = generatedDocId;
            } else {
                LastPart = part;
                id = LastNewIndexId;
            }

            row.SetPart(NJupiter::NPrivate::TIndexIdPart::Pack(id, part));
            if (OverrideRangeAndJob) {
                row.SetRange(0);
                row.SetJob(Max<ui32>() >> 1); // magic value to disable some asserts
            } else {
                row.SetJob(0);
            }
        }

        const google::protobuf::Descriptor* GetRowType() const override {
            return TMessage::descriptor();
        }

        THolder<google::protobuf::Message> CreateEmptyObject() const override {
            return MakeHolder<TMessage>();
        }

        ui32 GetDocId(const google::protobuf::Message&) const override {
            return Max<ui32>();
        }
    };

    TPortionsBundle::TPortionsBundle(NRtDoc::TBuilderLocalClientPtr client, TLog& log, TTimerLog& timerLog)
        : PrepClient(client)
        , Log(log)
        , TimerLog(timerLog)
    {
        Y_ENSURE(client);

        NRtDoc::TBuilderPortionsConfig::TPtr portionsConf = PrepClient->MutablePortionsConfig();
        PrepTypes = GetPrepTableTypes();

        // MergeCounts inputs: an explicit declaration as portions allows some of them to be missing
        if (portionsConf) {
            RegisterPortionIds(*portionsConf);
        }
    }

    // If you're editing set of processors, check out rtyserver_jupi/library/tables -- maybe it also needs in your fix
    void TPortionsBundle::RegisterPortionIds(NRtDoc::TBuilderPortionsConfig& cfg) {
        const auto tables = NJupiter::TShardsPrepareTables::Clientless();
        cfg.RegisterPortionId(TVector<TString>{
            NJupiter::GetYtPathName(tables.GetKeyInvCountsTable(0)),
            NJupiter::GetYtPathName(tables.GetAnnCountsTable(0)),
            NJupiter::GetYtPathName(tables.GetLinkAnnCountsTable(0)),
            NJupiter::GetYtPathName(tables.GetFactorAnnCountsTable(0)),
            NJupiter::GetYtPathName(tables.GetMetaTagsCountsTable(0)),
            NJupiter::GetYtPathName(tables.GetStatsTableLink()),
            NJupiter::GetYtPathName(tables.GetCountsKeyInvTable(0)),
            NJupiter::GetYtPathName(tables.GetPantherTable(0)),
        });
    }

    // If you're editing set of processors, check out rtyserver_jupi/library/tables -- maybe it also needs in your fix
    void TPortionsBundle::RegisterDocIdSetters(NRtDoc::ILocalClient& client) {
        auto tables = NJupiter::TShardsPrepareTables::Clientless();
        client.RegisterPreprocessor(MakeIntrusive<TWalrusDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TArcDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TCalcAttrsDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TContentAttrsDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TUrlHeavyDataDocId>());
        client.RegisterPreprocessor(MakeIntrusive<THostAttrsDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TExternalRelevAttrsDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TPantherDssmDocsBatchDocId<true>>(client.GetTotalInputDocCount()));
        client.RegisterPreprocessor(MakeIntrusive<TPantherDssmDocsBatchDocId<false>>(client.GetTotalInputDocCount()), NJupiter::GetYtPathName(tables.GetPantherDssmTermsTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<THnswDocEmbedDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TInvUrlHashDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TNavSrcStrictDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvEntryIndexId>());
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(), NJupiter::GetYtPathName(tables.GetPantherTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(), NJupiter::GetYtPathName(tables.GetCountsKeyInvTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(true, true), NJupiter::GetYtPathName(tables.GetNgramCountsTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(true, true), NJupiter::GetYtPathName(tables.GetPantherNgramsTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvOneDocId>(), NJupiter::GetYtPathName(tables.GetAttrsKeyInvTable(0)));
        client.RegisterPreprocessor(MakeIntrusive<TMobileUrlDocId>());
        client.RegisterPreprocessor(MakeIntrusive<TSimpleDocIdPreprocessor<NJupiter::TItdItpSlimIndexData>>());
    }

    // If you're editing set of processors, check out rtyserver_jupi/library/tables -- maybe it also needs in your fix
    void TPortionsBundle::AddDocIdSettersForStage(NRtDoc::ILocalClient& client, TPortionsBundle::EStage portionsStage) {
        const bool forceRenumCounts = true;
        const auto tables = NJupiter::TShardsPrepareTables::Clientless();

        if (portionsStage == EStage::MergePortions) {
            client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(forceRenumCounts),
                    NJupiter::GetYtPathName(tables.GetCountsKeyInvTable(0)));
            client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(forceRenumCounts, true),
                    NJupiter::GetYtPathName(tables.GetAttrsKeyInvTable(0)));
        }
        if (portionsStage == EStage::Final) {
            client.RegisterPreprocessor(MakeIntrusive<TOffroadKeyInvPortionIndexId>(false, false),
                    NJupiter::GetYtPathName(tables.GetAttrsKeyInvTable(0)));
        }
    }

    class TReducerRunner {
    protected:
        const TMercuryPrepCmd& Reducer;

        TVector<NRtDoc::TLocalJoinTable> Inputs;
        TVector<NRtDoc::TLocalJoinTable> Outputs;
        THolder<NYT::TTableReader<NProtoBuf::Message>> Reader;
        THolder<NYT::TTableWriter<NProtoBuf::Message>> Writer;
        TIntrusivePtr<NYT::IProtoReaderImpl> ReaderImpl;
    public:
        TReducerRunner(const TMercuryPrepCmd& reducer, const TPortionsBundle::TPrepTypesTable& prepTypes)
            : Reducer(reducer)
        {
            for (const TString& prepId : reducer.Method->GetInputs()) {
                if (!prepId) {
                    Inputs.push_back({nullptr, nullptr}); //IMercuryCmd::TSkipInput()
                    continue;
                }
                auto r = prepTypes.find(prepId);
                Y_ENSURE(r != prepTypes.end(), "no message type set for PrepId=" << prepId);
                const NProtoBuf::Descriptor* datatype = r->second;
                Inputs.push_back({prepId, datatype});
            }

            for (auto prepId : reducer.Method->GetOutputs()) {
                Outputs.push_back({prepId, nullptr});
            }
        }

        //note: non-virtual
        void Start(NRtDoc::ILocalJoinClient* client) {
            ReaderImpl = client->CreateLocalJoinReader(Inputs);
            Reader = MakeHolder<NYT::TTableReader<NProtoBuf::Message>>(ReaderImpl);
            Writer = MakeHolder<NYT::TTableWriter<NProtoBuf::Message>>(client->CreateLocalJoinWriter(Outputs));
        }

        virtual void Do() {
            THolder<NMercury::IProcAdapter> proc = THolder(Reducer.Method->Create());
            proc->Start(Writer.Get());
            do {
                proc->Do(Reader.Get(), Writer.Get());
                ReaderImpl->NextKey();
            } while (ReaderImpl->IsValid());
            proc->Finish(Writer.Get());
        }

        virtual void Finish() {
            Writer.Reset();
            Reader.Destroy();
            ReaderImpl.Drop();
        }
    };

    class TReducerTwoStageRunner : public TReducerRunner {
    public:
        using TReducerRunner::TReducerRunner;

        void Start(NRtDoc::ILocalJoinClient* client) {
            // creates SecondPassReader in advance, to avoid changes in PortionConfig
            SecondPassReader = client->CreateLocalJoinReader(Inputs);
            TReducerRunner::Start(client);
        }

        void Do() override {
            THolder<NMercury::IProcAdapter> proc = THolder(Reducer.Method->Create());
            constexpr ui32 nStages = 2;
            for (ui32 stage = 0; stage < nStages;) {
                proc->Start(Writer.Get());
                do {
                    proc->Do(Reader.Get(), Writer.Get());
                    ReaderImpl->NextKey();
                } while (ReaderImpl->IsValid());
                proc->Finish(Writer.Get());

                if (++stage < nStages) {
                    Restart();
                }
            }
        }

    private:
        void Restart() {
            Y_ENSURE(SecondPassReader);
            ReaderImpl = SecondPassReader;
            SecondPassReader.Drop();

            Reader = MakeHolder<NYT::TTableReader<NProtoBuf::Message>>(ReaderImpl);
        }
    private:
        TIntrusivePtr<NYT::IProtoReaderImpl> SecondPassReader;
    };

    class TRemapRunner: public TReducerRunner {
        // for some special cases - like when we need to write a portion, then convert it, then and replace the old with the converted.
    public:
        using TReducerRunner::TReducerRunner;

        void Start(NRtDoc::ILocalSegmentedClient* client) {
            Y_ENSURE(Inputs.size() == 1);
            Y_ENSURE(Outputs.size() == 1);
            Y_ENSURE(Inputs[0].PrepId == Outputs[0].PrepId);

            ReaderImpl = client->CreateLocalJoinReader(Inputs);
            Reader = MakeHolder<NYT::TTableReader<NProtoBuf::Message>>(ReaderImpl);
            Writer = MakeHolder<NYT::TTableWriter<NProtoBuf::Message>>(client->CreateLocalJoinWriter(Outputs, NRtDoc::TBuilderStage::NextGen()));
        }

        void Publish(NRtDoc::ILocalSegmentedClient* client) {
            // remove the old file, replace it with the new data
            client-> PublishMergedPortion(Outputs[0].PrepId, TFsPath());
        }
    };

    class TReducerMergeRunner : public TReducerRunner {
    protected:
        NRtDoc::ILocalSegmentedClient* SegClient;
        NRtDoc::TSegmentHitRemapper::TPtr Remapper;
    public:
        using TReducerRunner::TReducerRunner;

        void Start(NRtDoc::ILocalSegmentedClient* client, NRtDoc::TSegmentHitRemapper::TPtr remapper) {
            Y_ASSERT(client && remapper);

            SegClient = client;
            Remapper = remapper;

            Y_ASSERT(!ReaderImpl);
            Y_ASSERT(!Reader);
            Y_ENSURE(Inputs.size() == 1, "MergePortions() command is expected to have only one " << "input table");
            Y_ENSURE(Outputs.size() >= 1);
            Writer = MakeHolder<NYT::TTableWriter<NProtoBuf::Message>>(SegClient->CreateMergedPortionWriter(Outputs[0]));
        }

        void Do() override {
            Y_ASSERT(SegClient && Writer);

            const size_t segmentsCount = SegClient->GetSegmentsCount();
            const TVector<size_t> segmentIds = xrange<size_t>(segmentsCount);

            ReaderImpl = SegClient->CreatePrepPortionReader(Inputs[0], segmentIds, Remapper);
            Reader = MakeHolder<NYT::TTableReader<NProtoBuf::Message>>(ReaderImpl);

            THolder<NMercury::IProcAdapter> proc = THolder(Reducer.Method->Create());
            proc->Start(Writer.Get());
            proc->Do(Reader.Get(), Writer.Get());
            proc->Finish(Writer.Get());
        }

        void Finish() override {
            Y_ASSERT(SegClient);
            Writer.Reset();
        }

        void Publish(const TFsPath& outputDir, bool keepUnpublished) {
            if (keepUnpublished) {
                SegClient->PublishMergedPortion(Outputs[0].PrepId, outputDir);
            } else {
                SegClient->PublishMergedPortion(Outputs[0].PrepId, TString());
                SegClient->PublishSimplePortion(Outputs[0].PrepId, outputDir);
            }
        }
    };

    void TPortionsBundle::Run(const TMercuryCmds& allCommands) {
        NRtDoc::ILocalJoinClient* client = PrepClient.Get();

        for (const TMercuryPrepCmd& reducer : allCommands) {
            LOG << "PreparePortions " << reducer.DisplayName;
            TTimerWrapper w(TimerLog, "TPortionsBundle.Run.PreparePortions." + reducer.DisplayName);

            TReducerRunner runner(reducer, PrepTypes);
            runner.Start(client);
            runner.Do();
            runner.Finish();

            PrepClient->Finish(); //multiple calls of Finish() are allowed
        }
    }

    void TPortionsBundle::RunTwice(const TMercuryPrepCmd& twoStageCommand) {
        NRtDoc::ILocalJoinClient* client = PrepClient.Get();
        TTimerWrapper w(TimerLog, "TPortionsBundle.RunTwice.PreparePortions." + twoStageCommand.DisplayName);
        LOG << "PreparePortions " << twoStageCommand.DisplayName;
        TReducerTwoStageRunner runner(twoStageCommand, PrepTypes);
        runner.Start(client);
        runner.Do();
        runner.Finish();
        PrepClient->Finish(); //multiple calls of Finish() are allowed
    }


    void TPortionsBundle::MergePortions(const TMercuryPrepCmd& reducer, NRtDoc::TSegmentHitRemapper::TPtr remapper, const TString& publishDir, bool keepDelta) {
        TTimerWrapper w(TimerLog, "TPortionsBundle.MergePortions." + reducer.DisplayName);
        LOG << "MergePortions " << reducer.DisplayName;

        NRtDoc::ILocalSegmentedClient* client = PrepClient.Get();

        const size_t segmentsCount = client->GetSegmentsCount();
        for (size_t segment = 0; segment < segmentsCount; ++segment) {
            auto docIdMap = MakeHolder<NRtDoc::TDocIdMap>();
            client->LoadSegmentRemap(*docIdMap, segment);
            Y_ENSURE(docIdMap->Size() > 0);

            remapper->AddMapping(std::move(docIdMap), segment);
        }

        TReducerMergeRunner runner(reducer, PrepTypes);
        runner.Start(client, remapper);
        runner.Do();
        runner.Finish();

        PrepClient->Finish();
        if (publishDir) {
            runner.Publish(publishDir, keepDelta);
        }
    }

    void TPortionsBundle::ReduceMergedPortions(const TMercuryPrepCmd& reducer, const TString& /*publishDir*/) {
        TTimerWrapper w(TimerLog, "TPortionsBundle.ReduceMergedPortions." + reducer.DisplayName);
        LOG << "ReduceMergedPortions " << reducer.DisplayName;

        NRtDoc::ILocalJoinClient* client = PrepClient.Get();

        TReducerRunner runner(reducer, PrepTypes);
        runner.Start(client);
        runner.Do();
        runner.Finish();

        PrepClient->Finish();
    }

    void TPortionsBundle::RewritePortion(const TMercuryPrepCmd& reducer) {
        TTimerWrapper w(TimerLog, "TPortionsBundle.RewritePortion." + reducer.DisplayName);
        LOG << "RewritePortion " << reducer.DisplayName;

        NRtDoc::ILocalSegmentedClient* client = PrepClient.Get();

        TRemapRunner runner(reducer, PrepTypes);
        runner.Start(client);
        runner.Do();
        runner.Finish();
        PrepClient->Finish();
        runner.Publish(client);
    }
}
