#include "doc_extractor.h"
#include "merger_task.h"
#include "m2n_decoder.h"
#include "metrics.h"

#include <saas/rtyserver/config/merger_config.h>
#include <saas/rtyserver/common/fsync.h>
#include <saas/rtyserver/common/message_collect_server_info.h>
#include <saas/rtyserver/common/search_area_modifier.h>
#include <saas/rtyserver/common/should_stop.h>
#include <saas/rtyserver/components/generator/component.h>
#include <saas/rtyserver/indexer_core/file_index_modify_guard.h>
#include <saas/rtyserver/indexer_core/index_component_storage.h>
#include <saas/util/logging/exception_process.h>

namespace {
    class TMergerUpdatesStorageDecoder: public TDeferredUpdatesStorage::IDecoder {
    public:
        TMergerUpdatesStorageDecoder(IMergerTask& task, TFinalIndexPtr finalIndex, ui32 clusterId)
            : Task(task)
            , FinalIndex(finalIndex)
            , ClusterId(clusterId)
        {}

        bool Decode(const TDocSearchInfo::THash& identifier, const ui32 /*oldDocId*/, ui32& docId) const override {
            if (!FinalIndex->RemapIdentifier2DocId(identifier, docId)) {
                DEBUG_LOG << "Can't remap identifier " << identifier.Quote() << " into docid" << Endl;
                return false;
            }
            TRTYMerger::TAddress oldAddress;
            if (!Task.GetDecoder()->NewToOld(ClusterId, docId, oldAddress)) {
                DEBUG_LOG << "Can't remap identifier: " << identifier.Quote() << "/" << ClusterId << "/" << docId << " into oldDocId" << Endl;
                return false;
            }

            VERIFY_WITH_LOG(docId == Task.GetDecoder()->Decode(oldAddress.ClusterId, oldAddress.DocId).DocId,
                "Incorrect docid back mapping: %u, %u", oldAddress.DocId, docId);

            return true;
        }

    private:
        IMergerTask& Task;
        TFinalIndexPtr FinalIndex;
        ui32 ClusterId;
    };
}

const TVector<TString>& TMergeTask::GetDestSegments() const {
    return DestDirs;
}

void TMergeTask::MoveFromTemp(ui32 destIndex, IIndexStorage& storage, const std::atomic<bool>* rigidStop) {
    Y_UNUSED(storage);
    Y_UNUSED(rigidStop);
    CHECK_WITH_LOG(GetTempDestinations().size() == GetFullDestinations().size());
    CHECK_WITH_LOG(GetTempDestinations().size() > destIndex);
    const TFsPath tempDir(GetTempDestinations()[destIndex]);
    FsyncDirRecursively(tempDir);
    tempDir.RenameTo(GetFullDestinations()[destIndex]);
}

ui64 TMergeTask::Timestamp(ui32 clusterId, ui32 docId) const {
    CHECK_WITH_LOG(GetSourceSegments().size() > clusterId);
    CHECK_WITH_LOG(Storage);
    TFinalIndexPtr remapIndex = Storage->GetFinalIndex(GetSourceSegments()[clusterId]);
    CHECK_WITH_LOG(!!remapIndex);
    return remapIndex->GetDDKManager()->GetTimeLiveStart(docId);
}

void TMergeTask::DoBuildDecoder(IIndexStorage& storage) {
    Storage = dynamic_cast<TIndexStorage*>(&storage);
    NOTICE_LOG << "status=build_resort_map;" << Endl;
    TM2NDecoder::TOptions options;
    options.SegmentSize = Config.GetMergerConfig().MaxDocumentsToMerge;
    options.SizeDeviation = 0.15f;
    options.MaxDeadlineDocs = Config.GetMergerConfig().MaxDeadlineDocs;
    options.Pruning = Config.Pruning && Config.Pruning->PruningOn();
    options.ReuseDocids = options.Pruning || Config.GetMergerConfig().DocIdGenerator == NRTYServer::TMergerConfig::EDocIdGenerator::Default;
    options.PushSignals = !IsPortions;  // push signals only when NoRT segments are merging
    for (unsigned segment = 0; segment < GetSourceSegments().size(); ++segment) {
        NOTICE_LOG << "Build resort matrix for " << GetSourceSegments()[segment] << " (" << segment << "/" << GetSourceSegments().size() << ")" << Endl;
        TIndexControllerPtr remapIndex = storage.GetIndexController(GetSourceSegments()[segment]);
        VERIFY_WITH_LOG(!!remapIndex, "Incorrect merger task: segment %s doesn't exist", GetSourceSegments()[segment].c_str());
        Aggregators.push_back(MakeHolder<NRTYMerger::TAggregator>(remapIndex));
    }

    TM2NDecoder* decoder = new TM2NDecoder(options);
    Decoder.Reset(decoder);
    TDocumentsExtractor::TContext extractorContext(NSaas::Broadcast);
    TDocumentsExtractor docsExtractor(extractorContext);
    for (unsigned segment = 0; segment < GetSourceSegments().size(); ++segment) {
        NOTICE_LOG << "Build resort matrix for " << GetSourceSegments()[segment] << " (" << segment << "/" << GetSourceSegments().size() << ")" << Endl;
        TIndexControllerPtr remapIndex = storage.GetIndexController(GetSourceSegments()[segment]);
        TDocPlaceMap docPlaceMap;
        docsExtractor.GetDocPlaceMap(*remapIndex, docPlaceMap);
        decoder->AddInfo(remapIndex->GetDDKManager());
        ui32 currentTimeMinuntes = Now().Minutes();
        for (ui32 docid = 0; docid < docPlaceMap.size(); ++docid) {
            bool del = (Config.GetMergerConfig().ClearRemoved && remapIndex->IsRemoved(docid)) || !remapIndex->GetDDKManager()->CheckDeadLine(docid, currentTimeMinuntes);
            decoder->AddInfo(docid, del ? -1 : 0, docPlaceMap[docid], 0);
        }
    }
    decoder->Finalize();
    TString tempDirName, destDir;
    storage.AllocateIndex(destDir, tempDirName, GetShardNumber());
    DestDirs.push_back(destDir);
}

bool TMergeTask::Process(IMessage* message) {
    TMessageCollectServerInfo* infoMessage = dynamic_cast<TMessageCollectServerInfo*>(message);
    if (infoMessage) {
        infoMessage->AddMergerTaskInfo(GetJsonInfo());
        return true;
    }
    return false;
}

bool TMergeTask::DoOnStart(const std::atomic<bool>* /*rigidStop*/) {
    const TVector<TString>& tmpDests = GetTempDestinations();
    VERIFY_WITH_LOG(tmpDests.size() == 1, "Incorrect merger task");
    IndexMergeDir = tmpDests[0];
    return true;
}

void TMergeTask::DoOnBeforeMainStage() {
    ReportStart();
}

void TMergeTask::DoOnFailed() {
    ReportEnd(/*success=*/false);
    Notifier.OnMergerTaskFailed();

    ERROR_LOG << "Error on merge for " << IndexMergeDir << " directory" << Endl;
    if (!!UpdatesStorage && Storage) {
        for (TVector<TString>::const_iterator i = GetSourceSegments().begin(); i != GetSourceSegments().end(); ++i)
            Storage->GetFinalIndex(*i)->SetDeferredUpdatesStorage(nullptr);
    }
    for (ui32 i = 0; i < DestDirs.size(); ++i) {
        TString dir = DoBuildTempDest(*Storage, DestDirs[i]);
        if (NFs::Exists(dir)) {
            ERROR_LOG << "Directory " << dir << " removing" << Endl;
            RemoveDirWithContents(dir);
        }
        TString tmpfsDir = DoBuildTmpfsDest(*Storage, DestDirs[i]);
        if (NFs::Exists(tmpfsDir)) {
            ERROR_LOG << "Directory " << tmpfsDir << " removing" << Endl;
            RemoveDirWithContents(tmpfsDir);
        }
    }
    if (!!IndexMergeDir) {
        ERROR_LOG << "Merge failed. Deleting directory " << IndexMergeDir << " ..." << Endl;
        if (NFs::Exists(IndexMergeDir)) {
            RemoveDirWithContents(IndexMergeDir);
        }
        ERROR_LOG << "Merge failed. Deleting directory " << IndexMergeDir << " done" << Endl;
    }
}

void TMergeTask::DoOnFinished(const std::atomic<bool>* rigidStop)
try
{
    CHECK_WITH_LOG(Storage);
    THolder<TGuardTransaction> guardTransactionModifyFileIndexes;

    if (ShouldStop(rigidStop)) {
        ReportEnd(true);
        return;
    }

    const TBaseIndexComponent* comp = dynamic_cast<const TBaseIndexComponent*>(TIndexComponentsStorage::Instance().GetComponent(Config.IndexGenerator));
    CHECK_WITH_LOG(comp);
    bool indexIsEmpty = comp->IsEmptyIndex(GetFullDestinations()[0]);
    TFinalIndexGuardedPtr resultIndexGuard;
    TFinalIndexPtr resultIndex;

    if (!indexIsEmpty) {
        TString resultIndexDir(GetFullDestinations()[0]);

        AddProgressInfo("Add index...");
        resultIndexGuard = Storage->AddIndex(resultIndexDir, "", true, nullptr, NRTYServer::EExecutionContext::Merge);
        resultIndex = *resultIndexGuard;
        SendGlobalMessage<TUniversalAsyncMessage>("mergerDeletePhase");

        AddProgressInfo("Decoder building...");
        GetDecoder()->Decode(0, 0); // direct decoder building before indexation lock
        TRY // Restore actual indexfrq. Indexation wasn't locked...
            ui32 removedDocs = 0;
            for (ui32 att = 0; (att < 2); ++att) {
                if (att == 1) {
                    AddProgressInfo(ToString(GetShardNumber()) + " shard indexation locking...");
                    guardTransactionModifyFileIndexes.Reset(new TGuardTransaction(TFileIndexModifyGuard::Get()[GetShardNumber()]));
                    AddProgressInfo(ToString(GetShardNumber()) + " shard indexation locked");
                }
                if (indexIsEmpty)
                    continue;
                AddProgressInfo("Docids aggregation");
                TVector<ui32> docidsForDelete;
                TMap<ui32, TVector<ui32>> docidsForMark;
                for (ui32 segment = 0; segment < Segments.size(); ++segment) {
                    const auto agg = Aggregators[segment]->Detach();
                    for (auto&& i : agg.RemovedDocIds) {
                        ui32 docId = GetDecoder()->Decode(segment, i).DocId;
                        if (docId != REMAP_NOWHERE) {
                            docidsForDelete.push_back(docId);
                        }
                    }
                    for (auto&& marker : agg.MarkedDocIds) {
                        TVector<ui32>& docids = docidsForMark[marker.first];
                        for (auto&& i : marker.second) {
                            ui32 docId = GetDecoder()->Decode(segment, i).DocId;
                            if (docId != REMAP_NOWHERE) {
                                docids.push_back(docId);
                            }
                        }
                    }
                }
                removedDocs += docidsForDelete.size();
                AddProgressInfo("Docids aggregation OK(" + ToString(docidsForDelete.size()) + ")");
                if (!(indexIsEmpty = (GetDecoder()->GetSize() == removedDocs))) {
                    resultIndex->RemoveDocIds(docidsForDelete);
                    AddProgressInfo("Docids removed");
                    for (auto&& i : docidsForMark) {
                        resultIndex->MarkForDelete(i.second, i.first);
                    }
                    AddProgressInfo("Docids marked");
                } else {
                    AddProgressInfo("Index is empty");
                }
            }
            // make sure indexing is still locked and deactivate callbacks
            CHECK_WITH_LOG(guardTransactionModifyFileIndexes && guardTransactionModifyFileIndexes->WasAcquired());
            for (ui32 segment = 0; segment < Segments.size(); ++segment) {
                Aggregators[segment]->Deactivate();
            }
        CATCH_AND_RETHROW("while sync deleted.log");
    }
    Aggregators.clear();

    {
        IIndexStorage::TDeferredRemoveTaskPtr emptyRemove;
        THolder<TGuardOfTransactionSearchArea> hgtsa;
        TRY
            if (indexIsEmpty && !!resultIndex) {
                hgtsa.Reset(new TGuardOfTransactionSearchArea);
                AddProgressInfo("Merged index is empty: " + IndexMergeDir);
                resultIndex->Start();
                emptyRemove = Storage->RemoveIndex(resultIndex->Directory().PathName(), true);
                resultIndex.Reset(nullptr);
                resultIndexGuard.Reset(nullptr);
            }
        CATCH("ERROR after merging: remove empty result")
        if (!!emptyRemove) {
            emptyRemove->DestroyIndexes();
        }
    }

    AddProgressInfo("TGuardOfTransactionSearchArea: (shard " + ToString(GetShardNumber()) + ") locking...");
    {
        TGuardOfTransactionSearchArea gtsa;
        AddProgressInfo("TGuardOfTransactionSearchArea: (shard " + ToString(GetShardNumber()) + ") locked");
        try {
            if (!!resultIndex) {
                resultIndex->Start();
                DEBUG_LOG << "Index " << GetDestSegments()[0] << " started" << Endl;
            }
        } catch (...) {
            S_FAIL_LOG << "Can't start final index: " << CurrentExceptionMessage();
        }
        Storage->UnregisterSearchers(GetSourceSegments());
    }
    AddProgressInfo("TGuardOfTransactionSearchArea: (shard " + ToString(GetShardNumber()) + ") unlocked");

    IIndexStorage::TDeferredRemoveTaskPtr sourcesRemove = Storage->RemoveIndexes(GetSourceSegments(), true);
    AddProgressInfo("Source indexes removed");

    if (!!resultIndex && !!UpdatesStorage) {
        INFO_LOG << "Deferred updates applying... " << Endl;
        VERIFY_WITH_LOG(UpdatesStorage->IsEmpty() || GetDestSegments().size() == 1, "Indexing was not stopped before detach");
        UpdatesStorage->ApplyDeferredUpdates(*resultIndex->GetUpdater(), TMergerUpdatesStorageDecoder(*this, resultIndex, 0));
        INFO_LOG << "Deferred updates applying... OK" << Endl;
    }

    guardTransactionModifyFileIndexes.Reset(nullptr);
    AddProgressInfo("Source indexes removed");
    AddProgressInfo(ToString(GetShardNumber()) + " shard indexation unlocked");
    AddProgressInfo("Merging to " + IndexMergeDir + " finished");

    ReportEnd(true);
} catch (...) {
    ReportEnd(false);
    throw;
}

void TMergeTask::ReportStart() {
    if (GetIsPortions())
        return;

    CHECK_WITH_LOG(GetTaskStatus() == mtsActive);
    Metrics.MergeStartTime.Set(TInstant::Now().Seconds());
}
void TMergeTask::ReportEnd(bool success) {
    if (GetIsPortions())
        return;

    // no guarantee here that ReportStart() has been called.
    Metrics.MergeStartTime.Set(0);

    if (success) {
        if (GetTaskStatus() == mtsFinished) {
            Metrics.MergeCompleted.Inc();
            Metrics.MergeFailed.Set(0);
            Metrics.LastMergeTime.Set(TInstant::Now().Seconds());
        }
    } else {
        Metrics.MergeFailed.Inc();
        Metrics.MergeCompleted.Set(0);
    }
}
