#include "keyinv_op.h"

#include <robot/library/oxygen/indexer/processor/keyinv/makeaww.h>

#include <kernel/walrus/advmerger.h>

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

bool TKeyInvMerger::Merge(const std::atomic<bool>* stopKey, TMergeDocInfos& infos, const TRTYMerger::TContext& context) {
    {
        TSimpleSharedPtr<TAdvancedMergeTask> task(new TAdvancedMergeTask());
        task->StripKeys = context.StripKeys;
        task->UniqueHits = context.UniqueHits;
        task->TmpFileDir = context.TempDir;
        for (size_t i = 0; i < context.Sources.size(); ++i) {
            task->Inputs.push_back(TAdvancedMergeTask::TMergeInput(context.Sources[i] + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName));
        }
        for (size_t i = 0; i < context.Dests.size(); ++i) {
            task->Outputs.push_back(TAdvancedMergeTask::TMergeOutput(context.Dests[i] + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName));
        }
        infos.GenerateFinalRemapTable(task->FinalRemapTable);

        TAdvancedIndexMerger merger(task, stopKey);
        merger.SetExternalHitWriter(context.ExternalHitWriter);
        if (!merger.Run(context.WriteSentLens ? "sent" : nullptr)) {
            return false;
        }
    }

    //TODO:(yrum) osidorkin@ left this unfinished. The code below is never used
    if (context.MergeIndexAttrs) {
        TSimpleSharedPtr<TAdvancedMergeTask> task(new TAdvancedMergeTask());
        task->StripKeys = context.StripKeys;
        task->UniqueHits = context.UniqueHits;
        task->TmpFileDir = context.TempDir;
        for (size_t i = 0; i < context.Sources.size(); ++i) {
            const TString fname(TString::Join(
                context.Sources[i], ::INDEX_SUFFIX, "attrs", context.AdditionalSuffixIndexName));

            if (TFsPath(fname + "key").Exists()) {
                task->Inputs.push_back(TAdvancedMergeTask::TMergeInput(fname));
            }
        }
        if (!task->Inputs.empty()) {
            for (size_t i = 0; i < context.Dests.size(); ++i) {
                task->Outputs.push_back(TAdvancedMergeTask::TMergeOutput(context.Dests[i] + ::INDEX_SUFFIX + "attrs" + context.AdditionalSuffixIndexName));
            }
            infos.GenerateFinalRemapTable(task->FinalRemapTable);
            TAdvancedIndexMerger merger(task, stopKey);
            merger.SetExternalHitWriter(context.ExternalHitWriter);
            if (!merger.Run(nullptr)) {
                return false;
            }
        }
    }

    try {
        bool awwExists = false;
        for (auto&& source : context.Sources) {
            if (NFs::Exists(source + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName)) {
                awwExists = true;
                break;
            };
        }
        if (awwExists) {
            NOxygen::NKeyInv::TWordWeights weights;
            for (auto&& source : context.Sources) {
                NOxygen::NKeyInv::ReadAww(source + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName, &weights);
            }
            for (auto&& destination : context.Dests) {
                NOxygen::NKeyInv::MakeAww(destination + ::INDEX_SUFFIX + context.AdditionalSuffixIndexName, &weights);
            }
        }
    } catch (...) {
        WARNING_LOG << "aww merging failed: " << CurrentExceptionMessage() << Endl;
    }

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

    return true;
}

TString TKeyInvMerger::GetName() const {
    return "KeyInv";
}
