#include "builder.h"

#include "callbacks.h"
#include "parsed_entity.h"
#include "processor_profiler.h"

#include "processors/tuples_list.h"

#include <saas/rtyserver/config/const.h>
#include <saas/util/queue.h>

#include <robot/library/oxygen/indexer/processor/collection/collection.h>
#include <robot/library/oxygen/indexer/processor/hostiterator/hostiterator.h>
#include <robot/library/oxygen/indexer/processor/pruning/interface.h>
#include <robot/library/oxygen/indexer/object_context/object_context.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/protobuf/json/proto2json.h>

#include <util/string/join.h>

namespace {
    class TProcessorsBundles {
    public:
        TProcessorsBundles() {
            TSet<TString> bundleNames;
            Singleton<NOxygen::TBundleFactory>()->GetKeys(bundleNames);

            for (auto&& name : bundleNames) {
                INFO_LOG << "Creating Oxygen processors bundle " << name << Endl;
                THolder<NOxygen::IProcessorsBundle> bundle(NOxygen::TBundleFactory::Construct(name));
                VERIFY_WITH_LOG(bundle, "Cannot create bundle %s", name.data());
                Bundles.push_back(std::move(bundle));
            }
        }

        const TVector<THolder<NOxygen::IProcessorsBundle>>& GetBundles() const {
            return Bundles;
        }

    private:
        TVector<THolder<NOxygen::IProcessorsBundle>> Bundles;
    };

    const TVector<THolder<NOxygen::IProcessorsBundle>>& GetBundles() {
        return Singleton<TProcessorsBundles>()->GetBundles();
    }

    void WrapHostProcessors(TProcessorsBuilderContext& context) {
        Y_ENSURE(context.GroupAttrMap, "Can't wrap host processor without groupAttrMap");

        const NOxygen::TGroupAttrOptions& groupAttrOptions = context.OxygenConfig.GetGroupAttrOptions();
        NOxygen::TTupleNameSet groupAttrTupleNames;
        if (groupAttrOptions.GroupAttrTupleNamesSize()) {
            groupAttrTupleNames.insert(
                groupAttrOptions.GetGroupAttrTupleNames().begin(),
                groupAttrOptions.GetGroupAttrTupleNames().end()
            );
        } else if (groupAttrOptions.HasGroupAttrTupleName()) {
            groupAttrTupleNames.insert(groupAttrOptions.GetGroupAttrTupleName());
        } else {
            FAIL_LOG("Incorrect GroupAttr config");
        }

        NOxygen::TProcessorPtr collection = new NOxygen::TProcessorCollection(context.HostProcessors, context.InfoCallback);
        NOxygen::TProcessorPtr wrapper = new NOxygen::THostIterator(collection, *context.GroupAttrMap, groupAttrTupleNames);
        context.Processors.push_back(wrapper);
    }
}

TProcessorsBuilderContext::TProcessorsBuilderContext(const NOxygen::TOxygenOptions& oxygenConfig)
    : OxygenConfig(oxygenConfig)
{
    TaskPool = CreateRTYQueue(1, 0, NUtil::TSmartMtpQueue::TOptions::GetNoBallocOptions().SetThreadName("OxyBuilderCtx"));
}

TOXYIndexBuilder::TOXYIndexBuilder(
    const TFsPath directory,
    const NOxygen::TOxygenOptions& oxygenOptions,
    const TRTYServerConfig& globalConfig,
    bool pruneOnly,
    bool forMerge,
    const TTuplesListsAndFilters& tuplesListsAndFilters,
    const TString& componentName
)
    : TBaseIndexBuilder(globalConfig, componentName)
    , Directory(directory)
    , OxygenOptions(oxygenOptions)
    , ComponentConfig(*globalConfig.ComponentsConfig.Get<TRTYOxyConfig>(OXY_COMPONENT_NAME))
    , PruneOnly(pruneOnly)
    , ForMerge(forMerge)
    , TuplesListsAndFilters(tuplesListsAndFilters)
    , ProcessorProfiler(nullptr)
{
    auto composite = MakeHolder<NRTYServer::TCompositeCallback>();

    auto profiler  = MakeHolder<NRTYServer::TProcessorProfiler>();
    ProcessorProfiler = profiler.Get();
    composite->AddCallback(profiler.Release());

    CompositeCallback = composite.Release();
}

void TOXYIndexBuilder::BuildUpdater(TProcessorsBuilderContext& context, const NOxygen::TProcessorPtr& processor) {
    IOxyUpdater::TPtr updater = IOxyUpdater::TFactory::Construct(processor->GetClassName());
    INFO_LOG << "Updater for " << processor->GetClassName() << " building..." << Endl;
    if (!!updater && updater->Initialize(processor, context.OutputDirectory)) {
        INFO_LOG << "Updater for " << processor->GetClassName() << " built" << Endl;
        context.Updaters.push_back(updater);
    } else {
        INFO_LOG << "Updater for " << processor->GetClassName() << " building... CANCELLED" << Endl;
    }
}

bool TOXYIndexBuilder::BuildProcessors(TProcessorsBuilderContext& context) {
    try {
        auto localOxygenConfig = context.OxygenConfig;
        if (localOxygenConfig.HasStampTagOptions()) {
            if (!localOxygenConfig.GetStampTagOptions().HasDocTimestampTupleName()) {
                localOxygenConfig.MutableStampTagOptions()->SetTimestamp(Seconds());
            }
        }
        if (localOxygenConfig.HasShardConfOptions()) {
            localOxygenConfig.MutableShardConfOptions()->SetTimestamp(Seconds());
        }
        if (context.TuplesListsAndFilters) {
            context.Processors.push_back(new NOxygen::TTuplesListProcessor(context.OutputDirectory, *context.TuplesListsAndFilters));
        }

        for (auto&& bundle : GetBundles()) {
            bundle->CreateDocumentProcessors(context.Processors, localOxygenConfig, context);
            bundle->CreateHostProcessors(context.HostProcessors, localOxygenConfig, context);
        }

        if (!context.HostProcessors.empty()) {
            WrapHostProcessors(context);
        }

        if (context.ForUpdate) {
            context.Updaters.clear();
            for (ui32 i = 0; i < context.Processors.size(); ++i) {
                if (context.Processors[i]->IsWrapper()) {
                    auto wrapperProcessors = context.Processors[i]->GetProcessors();
                    if (wrapperProcessors) {
                        for (const auto& processor : *wrapperProcessors) {
                            BuildUpdater(context, processor);
                        }
                    }
                } else {
                    BuildUpdater(context, context.Processors[i]);
                }
            }
        }

        if (context.ForMerger) {
            context.MergersWithoutUpdaters.clear();
            context.MergersWithUpdaters.clear();

            TVector<NOxygen::TProcessorPtr> mergerProcessors;
            for (auto&& processor: context.Processors) {
                const TString& name = processor->GetClassName();

                INFO_LOG << "Merger for " << name << " building..." << Endl;
                IOxyMerger::TPtr merger = IOxyMerger::TFactory::Construct(name);
                if (!!merger && merger->Initialize(processor, context.MergerOptions.Get())) {
                    INFO_LOG << "Merger for " << name << " built" << Endl;
                    IOxyUpdater::TPtr updater = IOxyUpdater::TFactory::Construct(name);
                    if (!!updater) {
                        context.MergersWithUpdaters.push_back(merger);
                    } else {
                        context.MergersWithoutUpdaters.push_back(merger);
                    }
                } else {
                    INFO_LOG << "Merger for " << name << " building... CANCELLED" << Endl;
                    mergerProcessors.push_back(processor);
                }

                if (processor->IsWrapper() && name == "THostIterator" && processor->GetProcessors()) {
                    const NOxygen::TProcessors& collection = *processor->GetProcessors();
                    for (auto&& hostProcessor : collection) {
                        const TString& slaveName = hostProcessor->GetClassName();
                        IOxyMerger::TPtr merger = IOxyMerger::TFactory::Construct(slaveName);
                        if (!merger)
                            continue;
                        INFO_LOG << "Host data merger for " << slaveName << " building..." << Endl;
                        if (merger->Initialize(hostProcessor, context.MergerOptions.Get())) {
                            // we do not remove it from the mergerProcessors. Host merger is responsible itself for disabling the processor (may change in the future)
                            INFO_LOG << "Host data merger for " << slaveName << " built" << Endl;
                            context.MergersWithoutUpdaters.push_back(merger);
                        } else {
                            INFO_LOG << "Host data merger for " << slaveName << " building... CANCELLED" << Endl;
                        }
                    }
                }
            }
            context.Processors = std::move(mergerProcessors);
        }
        context.ProcessorsCollection = new NOxygen::TProcessorCollection(context.Processors, context.InfoCallback);
        return true;
    } catch (...) {
        ERROR_LOG << "Can't build processors: " << CurrentExceptionMessage() << Endl;
        return false;
    }
}

NOxygen::IPruningProcessorPtr TOXYIndexBuilder::BuildPruningProcessor(NOxygen::TProcessorPtr slave, const TProcessorsBuilderContext& context) {
    NOxygen::IPruningProcessorPtr result;
    for (auto&& bundle : GetBundles()) {
        NOxygen::IPruningProcessorPtr processor = bundle->CreatePruningProcessor(slave, context.OxygenConfig, context);
        if (result && processor) {
            ERROR_LOG << "Created two pruning processors: " << result->GetClassName() << " and " << processor->GetClassName() << Endl;
            AbortFromCorruptedConfig("Pruning configuration conflict");
        }
        if (processor) {
            result = processor;
        }
    }
    return result;
}

void TOXYIndexBuilder::InitMergerOptions(TOxyMergerConfig& config, const TRTYOxyConfig& componentConfig, const NOxygen::TOxygenOptions& /*oxyOptions*/) {
    const TRTYOxyConfig::TDbgOxyFlags& dbgFlags = componentConfig.GetDbgOxyFlags();
    if (dbgFlags.contains("regherf_external_merger")) {
        config.UseExternalRegHostErfMerger = true;
    }
}

bool TOXYIndexBuilder::Start() {
    try {
        TProcessorsBuilderContext context(OxygenOptions);
        context.SetDir(Directory).InitAttrs(AttrMap).SetForMerger(ForMerge).SetTuples(&TuplesListsAndFilters);
        context.InfoCallback = CompositeCallback;
        context.MergerOptions = MakeHolder<TOxyMergerConfig>();
        InitMergerOptions(*context.MergerOptions, ComponentConfig, OxygenOptions);
        if (!PruneOnly) {
            CHECK_WITH_LOG(BuildProcessors(context));
            CollectProcess = context.ProcessorsCollection;
        }
        PruningProcess = BuildPruningProcessor(CollectProcess, context);
        if (PruningProcess)
            CollectProcess = PruningProcess.Get();
        if (CollectProcess)
            CollectProcess->Start();
    } catch (...) {
        ERROR_LOG << "Start OXYIndexer failed: what = " + CurrentExceptionMessage() << Endl;
        return false;
    }
    return true;
}

bool TOXYIndexBuilder::DoClose(const NRTYServer::TBuilderCloseContext& context) {
    try {
        CHECK_WITH_LOG(PruneOnly || !!MapRemap || !DocsCount);
        if (CollectProcess) {
            CollectProcess->FinishEx(MapRemap.Get(), context.RigidStopSignal);
            CHECK_WITH_LOG(CollectProcess->GetVersion() == CollectProcess->GetVersion(Directory));
        }
        ReportStats();
    } catch (...) {
        ERROR_LOG << "Closing oxy indexer failed. message = " + CurrentExceptionMessage() << Endl;
        return false;
    }
    return true;
}

void TOXYIndexBuilder::CancelDocument(const TParsedDocument& document, const ui32 docId, const TString& message) {
    ERROR_LOG << message << " : " << document.GetDocSearchInfo().GetUrl() << Endl;
    FailedDocs.push_back(docId);
    document.SignalFailed();
}

bool TOXYIndexBuilder::IndexEmpty(int/* threadID*/, const ui32 /*docId*/) {
    CHECK_WITH_LOG(!PruningProcess);
    AtomicIncrement(DocsCount);
    return true;
}

void TOXYIndexBuilder::Index(int threadID, const TParsedDocument& document, const ui32 docId) {
    try {
        const TOXYParsedEntity* docIndex = document.GetComponentEntity<TOXYParsedEntity>(RTYConfig.IndexGenerator);
        VERIFY_WITH_LOG(docIndex, "there is no Index in doc %s", document.GetDocSearchInfo().GetUrl().data());
        VERIFY_WITH_LOG(!!docIndex->GetIndexedDoc(), "Incorrect service usage");

        NOxygen::TObjectContext oc(docIndex->GetIndexedDoc()->GetKiwiObject());
        auto check = TuplesListsAndFilters.GetRequiredTuples(ForMerge).VerifyObject(oc);
        if (!check.Correct) {
            TString message = "Incomplete KiwiObject: " + JoinSeq(" ", check.MissingTupleGroups);
            document.GetErrorsCollector().AddMessage<TLOG_INFO>(OXY_COMPONENT_NAME, message);
            CancelDocument(document, docId, message);

        } else if (!!CollectProcess) {
            NOxygen::TReturnObjectContext ret = CollectProcess->Process(oc, docId).GetValue(TDuration::Max());
            if (!ret.IsProcessedOk()) {
                document.GetErrorsCollector().AddMessage(OXY_COMPONENT_NAME, NProtobufJson::Proto2Json(ret.GetErrors()));
                CancelDocument(document, docId, "Document processing failed (" + NProtobufJson::Proto2Json(ret.GetErrors()) + ")");
            }
        }
    } catch (...) {
        CHECK_WITH_LOG(false) << "ProcessDoc failed: URL = " + document.GetDocSearchInfo().GetUrl() + " what = " + CurrentExceptionMessage() << Endl;
    }
    TBaseIndexBuilder::Index(threadID, document, docId);
}

void TOXYIndexBuilder::InitInteractions(const NRTYServer::IIndexBuildersStorage& storage) {
    FA = storage.GetBuilder<TBaseGeneratorBuilder>(FULL_ARCHIVE_COMPONENT_NAME);
    CHECK_WITH_LOG(FA);
}

ui32 TOXYIndexBuilder::GetResortMapForMerge(const IIndexController::TPtr index, TVector<ui32>& result) const {
    MapRemap.Reset(new NOxygen::TDocIdMap(DocsCount));
    if (!!PruningProcess) {
        for (ui32 i = 0; i < DocsCount; ++i) {
            if (index && index->IsRemoved(i)) {
                MapRemap->SetMapValue(i, NOxygen::TDocIdMap::DeletedDocument());
                DEBUG_LOG << "Removed DocId detected: " << i << Endl;
            }
        }
        DEBUG_LOG << "Docs for remap: " << MapRemap->GetSize() << "/" << DocsCount << Endl;
        *MapRemap = PruningProcess->CreateDocIdMapAsInFinish(MapRemap.Get());

//        VERIFY_WITH_LOG(MapRemap->GetSize() == DocsCount - FailedDocs.size(), "%d != %lu", MapRemap->GetSize(), DocsCount - FailedDocs.size());
        result.resize(DocsCount, -1);
        ui32 resultDocsCount = 0;

        for (auto&& i : FailedDocs) {
            CHECK_WITH_LOG(MapRemap->Map(i) == NOxygen::TDocIdMap::DeletedDocument());
        }

        for (ui32 i = 0; i < MapRemap->GetSize(); ++i) {
            if (MapRemap->Map(i) != MapRemap->DeletedDocument()) {
                result[i] = MapRemap->Map(i);
                ++resultDocsCount;
            }
            DEBUG_LOG << i << "->" << result[i] << Endl;
        }

        return resultDocsCount;
    } else {
        ui32 retRes = GetTrivialResortMap(index, result);
        for (ui32 i = 0; i < result.size(); ++i) {
            MapRemap->SetMapValue(i, result[i]);
        }
        return retRes;
    }
}

void TOXYIndexBuilder::ReportStats() {
    VERIFY_WITH_LOG(ProcessorProfiler, "Profiler is missing");
    auto timings = ProcessorProfiler->GetProcessingTimes();
    INFO_LOG << "Processor timings for " << Directory << "" << Endl;
    for (auto&& p : timings) {
        const NRTYServer::TProcessorProfiler::TProcessorTime& processor = p.second;
        INFO_LOG << processor.Name << ": Process " << processor.Process << ", Finish " << processor.Finish << Endl;
    }
}
