#include "multipart.h"
#include "multipart_ctx.h"

#include <kernel/multipart_archive/multipart.h>

using namespace NRTYArchive;

TString TMultipartArchiveMerger::GetName() const {
    return "MultipartArchive";
}

bool TMultipartArchiveMerger::Merge(const std::atomic<bool>* stopKey, TMergeDocInfos& /*infos*/, const TRTYMerger::TContext& context) {
    Y_ASSERT(context.MultipartMergerContext.Get());

    const TString fileName(!!context.MultipartArcFile ? ("/" + context.MultipartArcFile) : TString());

    for (ui32 d = 0; d < context.Dests.size(); ++d) {
        if (!context.Decoder->GetNewDocsCount(d))
            continue;
        TVector<NRTYArchive::TMultipartArchive::TRemap> remaps(context.Sources.size());
        for (ui32 docid = 0; docid < context.Decoder->GetNewDocsCount(d); ++docid) {
            if (!!stopKey && *stopKey) {
                return false;
            }

            TRTYMergerAddress oldAddr;
            if (!context.Decoder->NewToOld(d, docid, oldAddr))
                continue;
            remaps[oldAddr.ClusterId][docid] = oldAddr.DocId;
        }

        TArchiveOwner::TPtr dest(TArchiveOwner::Create(context.Dests[d] + fileName, context.MultipartMergerContext->Config,
                                                       context.Decoder->GetNewDocsCount(d)));
        for (ui32 s = 0; s < context.Sources.size(); ++s) {
            TArchiveOwner::TPtr multipart = TArchiveOwner::Create(context.Sources[s] + fileName, context.MultipartMergerContext->Config);
            if (multipart->IsWritable()) {
                multipart->Flush();
            }
            dest->Append(*multipart, &remaps[s]);
        }

        if (!context.MultipartMergerContext->DisablePartsOptimization) {
            const bool optimizeResult = dest->OptimizeParts(context.MultipartMergerContext->Config.CreateOptimizationOptions(), stopKey);
            if (!optimizeResult) {
                return false;
            }
        }
    }

    if (context.Callback) {
        context.Callback->OnAfterArcMerge();
    }

    return true;
}
