#include "document_modifier.h"
#include "index_updater.h"

#include <saas/library/behaviour/behaviour.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/common/message_search_doc_ids.h>
#include <saas/rtyserver/common/message_collect_server_info.h>
#include <saas/rtyserver/common/doc_utils.h>
#include <saas/rtyserver/common/search_area_modifier.h>

#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/common_indexers_config.h>

bool TDocumentModifier::Process(IMessage* message) {
    TMessageRegisterIndex* messRegIndex = dynamic_cast<TMessageRegisterIndex*>(message);
    if (messRegIndex) {
        TGuardTransaction g(Transaction);
        TString dirName = messRegIndex->GetIndexController()->Directory().BaseName();
        VERIFY_WITH_LOG(SegmentsMap.find(dirName) == SegmentsMap.end(), "Re-Register index for DocumentDeleter");

        SegmentsMap[dirName] = messRegIndex->GetIndexController();

        if (messRegIndex->GetIndexController()->GetType() & IIndexController::FINAL_MASK) {
            ui64 docsCount = messRegIndex->GetIndexController()->GetSearchableDocumentsCount();
            DocsCountPerSegment[dirName] = docsCount;
            SearchableDocumentsCount += messRegIndex->GetIndexController()->GetSearchableDocumentsCount();
        }

        INFO_LOG << "Directory " << dirName << " registered in DocumentModifier" << Endl;
        return true;
    }

    TMessageUnregisterIndex* messUnRegIndex = dynamic_cast<TMessageUnregisterIndex*>(message);
    if (messUnRegIndex) {
        TGuardTransaction g(Transaction);
        TString dirName = messUnRegIndex->GetIndexController()->Directory().BaseName();
        VERIFY_WITH_LOG(SegmentsMap.find(dirName) != SegmentsMap.end(), "Re-Unregister index for DocumentDeleter");
        SegmentsMap.erase(dirName);

        if (messUnRegIndex->GetIndexController()->GetType() & IIndexController::FINAL_MASK) {
            ui64 docsCount = DocsCountPerSegment[dirName];
            DocsCountPerSegment.erase(dirName);
            CHECK_WITH_LOG(docsCount <= SearchableDocumentsCount);
            SearchableDocumentsCount -= docsCount;
        }

        INFO_LOG << "Directory " << dirName << " unregistered from DocumentModifier" << Endl;
        return true;
    }

    if (TMessageModifyDocument* msgModify = dynamic_cast<TMessageModifyDocument*>(message)) {
        msgModify->Result = ModifyDocument(msgModify->Document, msgModify->Action);
        return true;
    }

    return false;
}

TModifyResult TDocumentModifier::ModifyDocumentByUrl(const TParsedDocument& document, NRTYServer::TMessage::TMessageType messageType, IModifyAction& action) {
    TGuardIncompatibleAction g(Transaction);
    TModifyResult result;

    const bool rejectDuplicates = Config.GetCommonIndexers().RejectDuplicates;
    const TDocSearchInfo& doc = document.GetDocSearchInfo();
    const ui64 version = document.GetVersion();
    const auto timestampCurrent = document.GetTimestamp();

    bool isRealUpdate = false;
    for (TSegmentsMap::const_iterator segment = SegmentsMap.begin(), e = SegmentsMap.end(); segment != e; ++segment) {
        IIndexController::TPtr segmentPtr = segment->second;

        TDocIdCandidate docIdCandidate(doc);
        if (segmentPtr->RemapUrl2DocIdCandidate(doc, docIdCandidate)) {
            ui32 docid = docIdCandidate.GetDocId();
            DEBUG_LOG << "Document " << doc.GetUrl() << " resolved in segment " << segment->first << " with docid " << docid << Endl;
            const IDDKManager* ddk = segmentPtr->GetDDKManager();
            if (ddk) {
                if (!docIdCandidate.IsVerified()) {
                    if (ddk->GetIdentifier(docid) != docIdCandidate.GetHash()) {
                        DEBUG_LOG << "Document " << doc.GetUrl() << " isn't contained in index, hash mismatch" << Endl;
                        continue;
                    }
                    docIdCandidate.SetVerified(true);
                }
                const ui64 previous = ddk->GetVersion(docid);
                const ui32 oldEntitiesHash = ddk->GetParsedEntitiesHash(docid);
                const auto timestampPrevious = ddk->GetTimeLiveStart(docid);
                if (IsDocOutdated(version, previous, timestampCurrent, timestampPrevious, Config.GetCommonIndexers().TimestampControlEnabled)) {
                    result.UpdateStatus(TModifyResult::OUTDATED);
                    result.Version = previous;
                    result.Timestamp = timestampPrevious;
                    continue;
                }

                if (rejectDuplicates && document.GetEntitiesHash() == oldEntitiesHash && GetBehaviour(messageType).IsDuplicatable && !segmentPtr->IsRemoved(docid)) {
                    result.UpdateStatus(TModifyResult::DUPLICATED);
                    continue;
                }
            }

            if (segmentPtr->IsRemoved(docid)) {
                DEBUG_LOG << "Document " << doc.GetUrl() << " removed in segment " << segment->first << " with docid " << docid << Endl;
                continue;
            }

            isRealUpdate = true;
            result.Count += action.ModifyDocumentByDocId(*segmentPtr, { docid });
        } else {
            DEBUG_LOG << "Document " << doc.GetUrl() << " not resolved in segment " << segment->first << Endl;
        }
    }
    if (!isRealUpdate && SearchableDocumentsCount >= Config.GetCommonIndexers().DocsCountLimit) {
        result.UpdateStatus(TModifyResult::REDUNDANT);
    }
    return result;
}

class TSearchInfoProcessorModify : public TMessageSearchDocIds::TSearchInfoProcessor {
public:
    struct THashHelper {
        inline size_t operator()(const IIndexController::TPtr& t) const noexcept {
            return (size_t)t.Get();
        }
    };
    typedef THashMap<IIndexController::TPtr, TVector<ui32>, THashHelper> TDocIdsRemoveMap;
private:
    TDocIdsRemoveMap DocIdsForRemove;
    TDocumentModifier::TSegmentsMap& SegmentsMap;
    IModifyAction& Action;
    ui64 TimeStart;
    ui32 ModifiedCount;
public:

    ui32 GetModifiedCount() const {
        return ModifiedCount;
    }

    bool IsEmpty() const {
        return !DocIdsForRemove.size();
    }

    TSearchInfoProcessorModify(TDocumentModifier::TSegmentsMap& segmentsMap, IModifyAction& action, NRTYServer::TTimestampValue /*actionTimestamp*/)
        : SegmentsMap(segmentsMap)
        , Action(action)
        , TimeStart(Now().MilliSeconds())
        , ModifiedCount(0)
    {
    }

    void Process(const TDocSearchInfo& info) override {
        const TString& indexName = info.GetIndexName();
        const ui32 docid = info.GetDocId();
        const TDocumentModifier::TSegmentsMap::const_iterator segment = SegmentsMap.find(indexName);
        VERIFY_WITH_LOG(segment != SegmentsMap.end(), "Incorrect segment in search result");
        DocIdsForRemove[segment->second].push_back(docid);
        TDocSearchInfo infoTemp("", 0);
        if (segment->second->DecodeIntoTempDocId(info, infoTemp)) {
            const TDocumentModifier::TSegmentsMap::const_iterator segmentTemp = SegmentsMap.find(infoTemp.GetIndexName());
            VERIFY_WITH_LOG(segmentTemp != SegmentsMap.end(), "Incorrect segment in search result");
            int docIdDecodedFromTemp = segmentTemp->second->DecodeFromTempDocId(infoTemp.GetDocId());
            if (docIdDecodedFromTemp != -1)
                DocIdsForRemove[segmentTemp->second].push_back(docIdDecodedFromTemp);
        }
    }

    int FinishPage(ui32 page) override {
        DEBUG_LOG << "action=remove_data;page=" << page << ";time=" << Now().MilliSeconds() - TimeStart << Endl;
        for (TSearchInfoProcessorModify::TDocIdsRemoveMap::const_iterator i = DocIdsForRemove.begin(), e = DocIdsForRemove.end(); i != e; ++i) {
            IIndexController::TPtr index = i->first;
            ui32 modCount = Action.ModifyDocumentByDocId(*index, i->second);

            if (i->first->HasSearcher())
                ModifiedCount += modCount;
            DEBUG_LOG << "action=clear_data;modif_count=" << modCount << ";page=" << page << ";index=" << i->first->Directory().BaseName() << ";time=" << Now().MilliSeconds() - TimeStart << Endl;
        }
        DEBUG_LOG << "action=remove_data_finished;page=" << page << ";time=" << Now().MilliSeconds() - TimeStart << Endl;
        DocIdsForRemove.clear();
        return -1;
    }

};

TModifyResult TDocumentModifier::ModifyDocumentByQuery(const TParsedDocument& document, IModifyAction& action)
{
    TGuardIncompatibleAction g(Transaction);
    TSearchInfoProcessorModify processor(SegmentsMap, action, document.GetTimestamp());
    const TString& query = document.GetQueryDel();
    TAtomicSharedPtr<TMessageSearchDocIds> messSearch = new TMessageSearchDocIds("text=" + query, &processor);
    INFO_LOG << "Query for delete: " << query << Endl;
    TInstant timess = Now();
    SendGlobalMessage(*messSearch);
    INFO_LOG << "Search: " << Now() - timess << Endl;

    TModifyResult result;
    result.Count = processor.GetModifiedCount();
    return result;
}

TModifyResult TDocumentModifier::ModifyDocument(const TQueuedDocument& document, IModifyAction& action) {

    const TParsedDocument& doc = document->GetDocument();
    {
        THolder<ISearchLocker::ILock> searchLock(SearchLocker->GetSearchLock());
        if (!doc.GetQueryDel()) {
            return ModifyDocumentByUrl(doc, document->GetCommand(), action);
        }

        const TString& query = doc.GetQueryDel();
        TModifyResult specialRemoveResult;
        if (query.StartsWith("$remove_kps$")) {
            TMessageRemoveSpecial messRemove(doc.GetDocSearchInfo().GetKeyPrefix());
            SendGlobalMessage(messRemove);
            specialRemoveResult.Count = messRemove.GetDocuments();
        } else if (query.StartsWith("$remove_all$")) {
            TMessageRemoveSpecial messRemove;
            SendGlobalMessage(messRemove);
            specialRemoveResult.Count = messRemove.GetDocuments();
        } else {
            return ModifyDocumentByQuery(doc, action);
        }

        return specialRemoveResult;
    }
}

TDocumentModifier::TDocumentModifier(const TRTYServerConfig& config, ISearchLocker* searchLocker)
    : Config(config)
    , SearchLocker(searchLocker)
{
    INFO_LOG << "Document modifier constructed" << Endl;
    RegisterGlobalMessageProcessor(this);
}

TDocumentModifier::~TDocumentModifier()
{
    INFO_LOG << "Document modifier destroyed" << Endl;
    UnregisterGlobalMessageProcessor(this);
}

TModifyActionVersionedDelete::TModifyActionVersionedDelete(TQueuedDocument document)
    : Document(document)
{
}

ui32 TModifyActionVersionedDelete::ModifyDocumentByDocId(IIndexController& index, const TVector<ui32>& docIds) {
    const TRTYMessageBehaviour& behaviour = GetBehaviour(Document->GetCommand());
    const TParsedDocument& parsedDocument = Document->GetDocument();
    const IDDKManager* ddk = index.GetDDKManager();

    if (behaviour.IsPureDeletion && ddk) {
        Y_ASSERT(parsedDocument.GetVersion() <= Max<ui32>());
        Y_ASSERT(parsedDocument.GetTimestamp() <= Max<ui32>());
        const auto version = static_cast<ui32>(parsedDocument.GetVersion());
        const auto timestamp = static_cast<ui32>(parsedDocument.GetTimestamp());
        for (auto&& docId : docIds) {
            ddk->UpdateVersionWithTimestamp(docId, version, timestamp);
        }
    }

    return index.RemoveDocIds(docIds);
}

TModifyActionMarkAsDeleted::TModifyActionMarkAsDeleted(const TString& indexDir)
    : Marker(FnvHash<ui32>(indexDir.data(), indexDir.size()))
{
    DEBUG_LOG << "Marker " << Marker << " built for " << indexDir << Endl;
}

ui32 TModifyActionMarkAsDeleted::ModifyDocumentByDocId(IIndexController& index, const TVector<ui32>& docIds) {
    return index.MarkForDelete(docIds, Marker);
}

TModifyActionUpdate::TModifyActionUpdate(TParsedDocument::TPtr document, const TString& deletionMarker)
    : DiffDocument(document)
    , DeletionMarker(deletionMarker)
    , Marker(FnvHash<ui32>(DeletionMarker.data(), DeletionMarker.size()))
{
    DEBUG_LOG << "Marker " << Marker << " made for " << DeletionMarker << Endl;
}

ui32 TModifyActionUpdate::ModifyDocumentByDocId(IIndexController& index, const TVector<ui32>& docIds) {
    IIndexUpdater* updater = index.GetUpdater();
    const bool isMemory = index.GetType() == IIndexController::MEMORY;
    ui32 result = 0;
    for (int i = 0; i < docIds.ysize(); ++i) {
        if (!DiffDocument->IsNeedRestore()) {
            result += (updater->UpdateDoc(docIds[i], DiffDocument) && !isMemory) ? 1 : 0;
        } else if (!isMemory) {
            TParsedDocument::TPtr restoredDoc;
            if (updater->RestoreDoc(docIds[i], restoredDoc)) {
                restoredDoc->ApplyPatch(*DiffDocument);
                DocumentsToIndex[restoredDoc->GetDocSearchInfo().GetHash()] = restoredDoc;
                ++result;
            }
        }
    }
    if (DiffDocument->IsNeedRestore()) {
        if (!DeletionMarker)
            index.RemoveDocIds(docIds);
        else
            index.MarkForDelete(docIds, Marker);
    }
    return result;
}

TMessageModifyDocument::TMessageModifyDocument(TQueuedDocument document, IModifyAction& action)
    : Document(document)
    , Action(action)
{}
