#include "index_builder.h"
#include "ext_indexer.h"
#include "rty_indexer.h"

#include <saas/rtyserver/config/common_indexers_config.h>
#include <saas/rtyserver/config/grouping_config.h>
#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/indexer_core/indexer_core.h>

#include <kernel/groupattrs/creator/mergeattrs.h>
#include <kernel/m2n/groupattrs.h>
#include <kernel/m2n/remaptable.h>
#include <kernel/tarc/merge/builder.h>

#include <kernel/keyinv/indexfile/indexstoragefactory.h>
#include <kernel/keyinv/indexfile/searchfile.h>

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

#include <yweb/config/pruning.h>

#include <ysite/yandex/erf/tf.h>


void TRTYIndexBuilder::RegisterStoragesAttr(TVector<IYndexStorage*> storages) {
    CHECK_EQ_WITH_LOG(storages.size(), BaseIndexers.size());
    for (ui32 i = 0; i < storages.size(); ++i)
        BaseIndexers[i]->RegisterStorageAttr(storages[i]);
}

void TRTYIndexBuilder::RegisterStoragesLemm(TVector<IYndexStorage*> storages) {
    CHECK_EQ_WITH_LOG(storages.size(), BaseIndexers.size());
    for (ui32 i = 0; i < storages.size(); ++i)
        BaseIndexers[i]->RegisterStorageLemm(storages[i]);
}

void TRTYIndexBuilder::RegisterInvCreator(TVector<NIndexerCore::IDirectTextCallback4*> callbacks) {
    CHECK_EQ_WITH_LOG(callbacks.size(), BaseIndexers.size());
    for (ui32 i = 0; i < callbacks.size(); ++i)
        BaseIndexers[i]->RegisterInvCreator(callbacks[i]);
}

void TRTYIndexBuilder::RegisterDTCallback(TVector<NIndexerCore::IDirectTextCallback2*> callbacks) {
    CHECK_EQ_WITH_LOG(callbacks.size(), BaseIndexers.size());
    for (ui32 i = 0; i < callbacks.size(); ++i)
        BaseIndexers[i]->RegisterDTCallback(callbacks[i]);
}

TRTYIndexBuilder::~TRTYIndexBuilder() {
    INFO_LOG << "status=rty_index_builder_destruction;temp=" << DirConfig.TempPrefix << Endl;
}

TRTYIndexBuilder::TRTYIndexBuilder(const NRTYServer::TIndexerConfig& config, const TDirConf& dirConfig,
    const TBaseConfig& baseConfig, const TString& componentName
    )
    : TBaseIndexBuilder(config.Common.Owner, componentName)
    , Config(config)
    , BaseConfig(baseConfig)
    , DirConfig(dirConfig)
    , MetaInfoCreator(new TRTYMetaInfoCreator)
    , SentenceLengthsWriter(new TRTYSentenceLengthWriter(dirConfig.NewPrefix, dirConfig.TempPrefix))
{
    ITextArchiveBuilder::TConstructContext builderCtx;
    builderCtx.DirConfig = DirConfig;
    builderCtx.DocsCount = BaseConfig.DocCount;
    builderCtx.ThreadsCount = Config.Threads;
    builderCtx.MultipartConfig = Config.Common.TextArchiveParams;
    TextArchiveWriter.Reset(ITextArchiveBuilder::TFactory::Construct(Config.Common.Owner.GetSearcherConfig().ArchiveType, builderCtx));

    TExternalIndexers ei;
    ei.MetaInfoCreator = MetaInfoCreator.Get();
    ei.SentenceLengthsWriter = SentenceLengthsWriter.Get();
    ei.TextArchiveWriter = TextArchiveWriter.Get();

    for (ui32 ithread = 0; ithread < Config.Threads; ithread++) {
        BaseIndexers.push_back(MakeHolder<TRTYIndexer>(dirConfig, baseConfig, ithread, Config.Common.Owner, ei));
    }

}

void TRTYIndexBuilder::Index(int threadID, const TParsedDocument& document, const ui32 docId) {
    CHECK_WITH_LOG((ui32)threadID < BaseIndexers.size());
    try {
        BaseIndexers[threadID]->IndexDoc(document);
        BaseIndexers[threadID]->IncDoc(docId);
    } catch (yexception& e) {
        CHECK_WITH_LOG(false) << "ProcessDoc failed: URL = " + document.GetDocSearchInfo().GetUrl() + " what = " + TString(e.what()) << Endl;
    }
    TBaseIndexBuilder::Index(threadID, document, docId);
}

namespace {

TString GetTempPortionPrefix(const TString& tempprefix, const ui32 portionNum) {
    return TString::Join(tempprefix, ToString(portionNum), '_');
}

void MergePortions(const TString& tempdir, const TString& tempprefix, const TString& newprefix,
    const TVector<ui32>& remap, NGroupingAttrs::TVersion grattrVersion, const ui32 portionsCount)
{
    TSimpleSharedPtr<TAdvancedMergeTask> task(new TAdvancedMergeTask());
    task->UniqueHits = true;
    task->StripKeys = false;
    task->FinalRemapTable.Create(1, 0); // 1 - single cluster, 0 - append documents by default

    NIndexerCore::TIndexStorageFactory walrus;
    walrus.InitIndexResources(tempprefix.data(), tempdir.data(), nullptr);

    TVector<TString> newAttrNames;
    TVector<ui32> docs;
    for (ui32 i = 0; i < portionsCount; i++) {
        const TString newtempprefix = GetTempPortionPrefix(tempprefix, i);
        walrus.LoadYandexPortions((newtempprefix + "pls").data());
        newAttrNames.push_back(newtempprefix + "aa");
        TFileInput f(newtempprefix + "plsad");
        ::Load(&f, docs);
        for (size_t j = 0; j < docs.size(); j++) {
            task->FinalRemapTable.SetOutputCluster(0, docs[j], 0);
        }
        docs.clear();
    }
    if (!remap.empty()) {
        for (ui32 id = 0; id < remap.size(); ++id) {
            if (remap[id] == (ui32)-1)
                task->FinalRemapTable.SetDeletedDoc(0, id);
            else
                task->FinalRemapTable.SetNewDocId(0, id, remap[id]);
        }
    }
    AddInputPortions(*task, walrus.GetNames());
    task->Outputs.push_back(TAdvancedMergeTask::TMergeOutput(walrus.GetIndexPrefix()));
    {
        TAdvancedIndexMerger merger(task);
        merger.Run();
    }
    walrus.RemovePortions();

    TString attrTempFile(tempprefix + "aa");
    NGroupingAttrs::TConfig::Mode grattrMode = NGroupingAttrs::TConfig::Search;
    NGroupingAttrs::MergeAttrsPortion(newAttrNames, nullptr,
        attrTempFile, nullptr, grattrMode, grattrVersion);
    if (!remap.empty() && NFs::Exists(attrTempFile)) {
        NM2N::TCompactMultipleRemapTable remapTable;
        remapTable.UpdateMaxDstClusterAndDocId(0, remap.size() - 1);
        remapTable.ResizeHolders(TVector<ui32>(1, remap.size()));
        for (size_t i = 0; i < remap.size(); ++i) {
            if (remap[i] != (ui32)-1)
                remapTable.AddRemap(i, 0, remap[i], 0);
        }
        const TString outputFile(tempprefix + "remappedaa");
        {
            const char* p1 = attrTempFile.c_str();
            const char* p2 = outputFile.c_str();
            NM2N::TAttrsRemap remap2(&p1, &p2, &remapTable, false, grattrVersion);
            remap2.Remap();
            remap2.WriteResult();
        }
        remove(attrTempFile.c_str());
        attrTempFile = outputFile;
    }

    for (ui32 j = 0; j < portionsCount; j++) {
        const TString newtempprefix = GetTempPortionPrefix(tempprefix, j);
        remove((newtempprefix + "aa").data());
        remove((newtempprefix + "pls").data());
        remove((newtempprefix + "plsad").data());
    }
    remove((tempprefix + "p").data());

    rename((tempprefix + "key").data(), (newprefix + "key").data());
    rename((tempprefix + "inv").data(), (newprefix + "inv").data());

    if (NFs::Exists(attrTempFile))
        rename(attrTempFile.data(), (newprefix + "aa").data());

    TYndex4Searching yndex;
    yndex.InitSearch(newprefix);
    SaveYandLengths(yndex, newprefix.data(), /*nofrq*/ false, /*check*/ false, /*notf*/ true, /*noerf*/ true);
}

}

bool TRTYIndexBuilder::DoClose(const NRTYServer::TBuilderCloseContext& context) {
    if (!!context.RigidStopSignal && *context.RigidStopSignal)
        return true;
    VERIFY_WITH_LOG(context.SrcDir == context.DstDir, "Not supported");
    VERIFY_WITH_LOG(context.RemapTable, "Incorrect remap table usage for IndexerBuilder");

    MetaInfoCreator->SaveC2NImpl(DirConfig.NewPrefix);
    MetaInfoCreator.Reset(nullptr);

    SentenceLengthsWriter->Close(context.RemapTable, context.RigidStopSignal);
    SentenceLengthsWriter.Reset(nullptr);

    TextArchiveWriter->Close(context.RemapTable);

    NGroupingAttrs::TVersion version = Config.Common.GroupingConfig->GetVersion();
    MergePortions(DirConfig.TempDir,
                  DirConfig.TempPrefix,
                  DirConfig.NewPrefix,
                  *context.RemapTable,
                  version,
                  Config.Threads);
    return true;
}

bool TRTYIndexBuilder::Start() {

    DEBUG_LOG << "status=rty_index_builder_start;temp=" << DirConfig.TempPrefix << ";threads=" << Config.Threads << Endl;
    for (ui32 ithread = 0; ithread < Config.Threads; ithread++) {
        BaseIndexers[ithread]->Start();
    }
    DEBUG_LOG << "status=rty_indexers_started;temp=" << DirConfig.TempPrefix << Endl;
    TextArchiveWriter->Start();
    DEBUG_LOG << "status=rty_index_builder_started;temp=" << DirConfig.TempPrefix << Endl;
    return true;
}

bool TRTYIndexBuilder::Stop() {
    VERIFY_WITH_LOG(BaseIndexers.size() == Config.Threads, "Incorrect Stop call");
    for (ui32 ithread = 0; ithread < BaseIndexers.size(); ithread++) {
        BaseIndexers[ithread]->Stop();
    }
    BaseIndexers.clear();
    TextArchiveWriter->Stop();
    return true;
}
