#include "archive_op.h"

#include <saas/rtyserver/common/should_stop.h>

#include <library/cpp/streams/special/buffered_throttled_file.h>

#include <ysite/yandex/srchmngr/arcmgr.h>

#include <kernel/tarc/iface/tarcio.h>

#include <util/folder/path.h>
#include <util/generic/ptr.h>
#include <util/datetime/base.h>

bool TArchiveMerger::Merge(const std::atomic<bool>* stopKey, TMergeDocInfos& /*infos*/, const TRTYMerger::TContext& context) {
    TVector<THolder<TArchiveIterator>> inputs;
    TArchiveVersion archiveVersion = 0;
    size_t nSourcesUsed = 0;
    TString indexName = ::INDEX_SUFFIX + context.AdditionalSuffixIndexName + "arc";
    for (ui32 i = 0; i < context.Sources.size(); ++i) {
        TFsPath fs(context.Sources[i] + "/" + indexName);
        if (!fs.Exists()) {
            inputs.push_back(THolder<TArchiveIterator>());
            continue;
        }

        inputs.push_back(MakeHolder<TArchiveIterator>());
        inputs.back()->Open(fs.Fix().c_str());

        const TArchiveVersion version = inputs.back()->GetArchiveVersion();
        if (!archiveVersion) {
            archiveVersion = version;
        }
        if (archiveVersion != version) {
            ERROR_LOG << "Cannot merge mixed-versioned archives" << Endl;
            return false;
        }
        nSourcesUsed++;
    }

    if (nSourcesUsed < context.Sources.size()) {
        if (nSourcesUsed == 0) {
            ERROR_LOG << indexName << " not found, cannot merge" << Endl;
            return true; //index is not created
        } else {
            ERROR_LOG << indexName << " is not found in some segments" << Endl;
        }
    }

    TVector<TSimpleSharedPtr<TBufferedThrottledFileOutputStream> > archOut;
    TVector<TVector<ui64> > dirData;

    for (unsigned i = 0; i < context.Dests.size(); i++) {
        archOut.push_back(new TBufferedThrottledFileOutputStream(context.Dests[i] + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName + "arc", context.WriteOption));
        WriteTextArchiveHeader(*archOut.back(), archiveVersion);
    }
    dirData.resize(context.Dests.size());

    TVector<ui64> offset(context.Dests.size(), ARCHIVE_FILE_HEADER_SIZE);
    for (ui32 clFrom = 0; clFrom < context.Sources.size(); clFrom++) {
        if (ShouldStop(stopKey)) {
            DEBUG_LOG << "Archive merger op interrupted by stop signal" << Endl;
            return false;
        }
        if (!inputs[clFrom])
            continue;

        TArchiveIterator& iterator = *inputs[clFrom];
        TArchiveHeader* header = nullptr;

        ui32 clusterSize = context.Decoder->GetSizeOfCluster(clFrom);
        for (ui32 docFrom = 0; ((header = iterator.NextAuto()) != nullptr) && clusterSize; ++docFrom) {
            if (ShouldStop(stopKey)) {
                DEBUG_LOG << "Archive merger op interrupted by stop signal" << Endl;
                return false;
            }
            clusterSize--;
            if (!context.Decoder->Check(clFrom, header->DocId))
                continue;
            TRTYMerger::TAddress addr = context.Decoder->Decode(clFrom, header->DocId);
            if (addr.IsRemoved())
                continue;
            auto& dir = dirData[addr.ClusterId];
            header->DocId = addr.DocId;
            archOut[addr.ClusterId]->Write(header, header->DocLen);
            if (dir.size() <= addr.DocId) {
                dir.resize(addr.DocId + 1, -1);
            }
            dir[addr.DocId] = offset[addr.ClusterId];
            offset[addr.ClusterId] += header->DocLen;
        }
    }

    archOut.clear();
    inputs.clear();

    for (unsigned i = 0; i < context.Dests.size(); i++) {
        TFile f(context.Dests[i] + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName + "dir", CreateAlways | WrOnly | Seq | Direct);
        if (dirData[i].size())
            f.Write(&dirData[i][0], dirData[i].size() * sizeof(dirData[i][0]));
    }

    if (context.Callback)
        context.Callback->OnAfterArcMerge();
    return true;
}
