#include "memory_manager.h"

#include "light_entity.h"
#include "light_layers_registry.h"
#include "parsed_entity.h"

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

TMemoryFAManager::TMemoryFAManager(const TRTYServerConfig& config)
    : IFAManager(config)
    , Storage(config.GetIndexerMemoryConfig().MaxDocuments)
{
}

bool TMemoryFAManager::UpdateDoc(ui32 docId, const TParsedDocument::TPtr doc) {
    if (Y_UNLIKELY(!IsOpened()))
        return false;
    TParsedDocument::TPtr oldDoc = GetDoc(docId, NRTYServer::TMessage::ADD_DOCUMENT, true);
    if (Y_UNLIKELY(!oldDoc))
        return false;
    oldDoc->ApplyPatch(*doc);
    Index(*oldDoc, docId);
    return true;
}

ui32 TMemoryFAManager::DoRemoveDocids(const TVector<ui32>& docids) {
    ui32 result = 0;
    for (auto&& i : docids) {
        auto& r = Storage[i];
        TGuard g(r.Lock);
        TRecord emptyRecord;
        emptyRecord.Swap(r);
        g.Release();
        if (emptyRecord.Doc) {
            ++result;
            DEBUG_LOG << "Remove docid " << i << " from FAMemory" << Endl;
        }
        else {
            DEBUG_LOG << "Can't remove docid " << i << " from FAMemory" << Endl;
        }
    }
    return result;
}

bool TMemoryFAManager::IsRemoved(ui32 docid) const {
    auto& r = Storage[docid];
    TGuard g(r.Lock);
    return !r.Doc;
}

void TMemoryFAManager::Remap(const TVector<ui32>* remapTable) {
    CHECK_WITH_LOG(!remapTable);
}

ui32 TMemoryFAManager::GetDocumentsCountImpl() const {
    return Storage.size();
}

TParsedDocument::TPtr TMemoryFAManager::GetDoc(ui32 docId, NRTYServer::TMessage::TMessageType command, bool unsafe) const {
    Y_UNUSED(command);
    Y_UNUSED(unsafe);
    auto& r = Storage[docId];
    TGuard g(r.Lock);
    return r.Doc;
}

TBlob TMemoryFAManager::ReadRawDoc(const TString& layer, ui32 docId) const {
    auto& r = Storage[docId];
    if (TLightLayersRegistry::Instance()->IsLightLayer(layer)) {
        TGuard g(r.Lock);
        if (auto e = MapFindPtr(r.LightEntities, layer)) {
            return *e;
        }
    } else {
        TGuard g(r.Lock);
        if (auto doc = r.Doc) {
            g.Release();
            NRTYServer::TDocSerializeContext serializeContext(layer);
            NRTYServer::TParsedDoc pd;
            doc->MergeToProto(pd, serializeContext);
            TBufferOutput out;
            ::google::protobuf::io::TProtoSerializer::Save(&out, pd);
            return TBlob::FromBuffer(out.Buffer());
        }
    }
    return {};
}

bool TMemoryFAManager::SearchDocument(ui32 docId, const TCgiParameters& cgi, NMetaProtocol::TDocument& doc, const TFAExtraSearchParams* extraParams) const {
    auto parsedDocument = GetDoc(docId);
    if (!parsedDocument) {
        return false;
    }
    auto entity = parsedDocument->GetComponentEntity<TRTYFullArchiveParsedEntity>(FULL_ARCHIVE_COMPONENT_NAME);
    if (!entity) {
        return false;
    }
    CopyDocumentToReport(docId, &entity->GetDocument(), cgi, doc, extraParams);
    return true;
}

bool TMemoryFAManager::SearchDocumentCandidate(TDocIdCandidate docId, const TCgiParameters& cgi, NMetaProtocol::TDocument& doc, const TFAExtraSearchParams* extraParams) const {
    auto parsedDocument = GetDoc(docId.GetDocId());
    if (!parsedDocument) {
        return false;
    }
    if (!docId.IsVerified() && parsedDocument->GetDocSearchInfo().GetHash() != docId.GetHash()) {
        return false;
    }
    auto entity = parsedDocument->GetComponentEntity<TRTYFullArchiveParsedEntity>(FULL_ARCHIVE_COMPONENT_NAME);
    if (!entity) {
        return false;
    }
    CopyDocumentToReport(docId.GetDocId(), &entity->GetDocument(), cgi, doc, extraParams);
    return true;
}

void TMemoryFAManager::Index(const TParsedDocument& document, const ui32 docId) {
    NRTYServer::TDocSerializeContext context;
    NRTYServer::TParsedDoc pd;
    document.MergeToProto(pd, context);
    TRecord newRecord;
    newRecord.Doc = MakeAtomicShared<TParsedDocument>(GetRTYConfig());
    newRecord.Doc->FillFromProto(pd, NRTYServer::TMessage::ADD_DOCUMENT, NRTYServer::TDocParseContext());
    for (const auto& layer : TLightLayersRegistry::GetRegisteredComponents()) {
        const auto* entity = newRecord.Doc->GetComponentEntity<TRTYFullArchiveLightEntity>(layer);
        if (entity && !entity->GetData().IsNull()) {
            newRecord.LightEntities[layer] = entity->GetData();
        }
    }
    auto& r = Storage[docId];
    TGuard g(r.Lock);
    r.Swap(newRecord);
}

bool TMemoryFAManager::GetDocInfo(const ui32 docId, NJson::TJsonValue& result) const {
    auto document = GetDoc(docId);
    if (!document) {
        return false;
    }

    NRTYServer::TParsedDoc pd;
    NRTYServer::TDocSerializeContext context;
    document->MergeToProto(pd, context);
    return IFAManager::GetDocInfo("memory", pd, result);
}

void TMemoryFAManager::TRecord::Swap(TRecord& other) {
    Doc.Swap(other.Doc);
    LightEntities.swap(other.LightEntities);
}
