#include "disk_manager.h"
#include "light_layers_registry.h"
#include "config.h"

#include <saas/util/logging/exception_process.h>

namespace {
    TString GetLegacyArchiveName(const TString& layer) {
        return OLD_FULL_ARC_FILE_NAME_PREFIX + layer;
    }
}

TDiskFAManager::TDiskFAManager(
    const TFsPath& dir,
    ui32 docCount,
    const TRTYServerConfig& config,
    ui32 writeSpeedBytes,
    const TSet<TString>& layers,
    bool enableFastUpdates,
    bool readOnly,
    bool isFinal)
    : IFAManager(config)
    , Directory(dir)
    , ReadOnly(readOnly)
    , FastUpdatesProcessing(enableFastUpdates)
    , DocsCount(docCount)
{
    if (!isFinal) {
        RemapMutex.ConstructInPlace();
    }

    const TRTYFullArchiveConfig* faConfig = Config.ComponentsConfig.Get<TRTYFullArchiveConfig>(FULL_ARCHIVE_COMPONENT_NAME);
    CHECK_WITH_LOG(faConfig);

    for (const auto& layer : layers) {
        AddLayer(Config, faConfig->GetLayer(layer), dir, layer, writeSpeedBytes, ReadOnly);
    }
    if (layers.empty()) {
        TIndexMetadataProcessor metaData(Directory);
        for (ui32 i = 0; i < metaData->GetFullArcHeader().WritedLayersSize(); ++i) {
            const TString& layer = metaData->GetFullArcHeader().GetWritedLayers(i);
            AddLayer(Config, faConfig->GetLayer(layer), dir, layer, writeSpeedBytes, ReadOnly);
        }
    }
    if (Layers.empty()) {
        ythrow yexception() << "cannot get layerList " << Directory;
    }
    TLayers::const_iterator i = Layers.find(NRTYServer::NFullArchive::FullLayer);
    if (i != Layers.end())
        FullLayer = i->second;
    MainLayer = faConfig->GetMainLayer();
}

void TDiskFAManager::AddLayer(const TRTYServerConfig& config, const NRTYArchive::TMultipartConfig& lc, const TString& directory, const TString& layerName, ui32 writeSpeedBytes, bool readOnly) {
    if (TLightLayersRegistry::Instance()->IsLightLayer(layerName)) {
        Layers.insert(std::make_pair(layerName, new TFALightLayer(config, lc, directory, layerName, writeSpeedBytes, readOnly)));
    } else {
        Layers.insert(std::make_pair(layerName, new TFAKVLayer(Config, lc, directory, layerName, writeSpeedBytes, readOnly)));
    }
}

bool TDiskFAManager::GetDocInfo(const ui32 docId, NJson::TJsonValue& result) const {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    bool success = false;
    for (const auto& layer : Layers) {
        NRTYServer::TParsedDoc pd;
        if (!layer.second->ReadParsedDoc(pd, docId, nullptr))
            continue;
        success |= IFAManager::GetDocInfo(layer.first, pd, result);
    }
    return success;
}

ui32 TDiskFAManager::GetDocumentsCountImpl() const {
    return DocsCount;
}

ui64 TDiskFAManager::GetLockedMemorySize() const {
    ui64 size = 0;
    for (auto& layer : Layers)
        size += layer.second->GetLockedMemorySize();
    return size;
}

bool TDiskFAManager::UpdateDoc(ui32 docId, const TParsedDocument::TPtr doc) {
    if (!FastUpdatesProcessing)
        return true;
    ui32 newDocId = docId;
    if (Y_UNLIKELY(!IsOpened()))
        return false;
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    if (Y_UNLIKELY(!!RemapTable))
        newDocId = RemapTable->at(docId);
    TParsedDocument::TPtr oldDoc = GetDoc(newDocId, NRTYServer::TMessage::ADD_DOCUMENT, true);
    if (Y_UNLIKELY(!oldDoc))
        return false;
    oldDoc->ApplyPatch(*doc);
    IndexUnsafe(*oldDoc, newDocId);
    return true;
}

ui32 TDiskFAManager::DoRemoveDocids(const TVector<ui32>& docids) {
    CHECK_WITH_LOG(!ReadOnly);

    TTryReadGuardBase<TMaybe<TRWMutex>> remapGuard(RemapMutex);
    if (!remapGuard.WasAcquired()) {
        for (auto&& i : docids) {
            DEBUG_LOG << "Add deferred remove command for " << Directory.GetPath() << "/" << i << Endl;
        }
        TGuard<TMutex> g(DeferredRemovesMutex);
        DeferredRemoves.insert(DeferredRemoves.end(), docids.begin(), docids.end());
        return 0;
    }

    return DoRemoveDocidsUnsafe(docids);
}

ui32 TDiskFAManager::DoRemoveDocidsUnsafe(const TVector<ui32>& docids) {
    ui32 result = 0;
    for (ui32 docid : docids) {
        bool success = false;
        ui32 newDocId = !!RemapTable ? RemapTable->at(docid) : docid;
        CHECK_WITH_LOG(newDocId != REMAP_NOWHERE);
        for (auto& layer : Layers) {
            if (layer.second->RemoveDocument(newDocId))
                success = true;
        }
        if (success) {
            DEBUG_LOG << "Remove command for " << Directory.GetPath() << "/" << docid << Endl;
            ++result;
        }
    }
    return result;
}

void TDiskFAManager::Index(const TParsedDocument& document, const ui32 docId) {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    IndexUnsafe(document, docId);
}

void TDiskFAManager::IndexUnsafe(const TParsedDocument& document, const ui32 docId) {
    CHECK_WITH_LOG(!ReadOnly);
    for (auto& layer : Layers)
        layer.second->Index(document, docId);
}

void TDiskFAManager::IndexUnsafe(const NRTYServer::TParsedDoc& document, const ui32 docId) {
    CHECK_WITH_LOG(!ReadOnly);
    for (auto& layer : Layers)
        layer.second->Index(document, docId);
}

void TDiskFAManager::IndexUnsafe(const TBlob& document, const ui32 docId) {
    CHECK_WITH_LOG(!ReadOnly);
    for (auto& layer : Layers)
        layer.second->Index(document, docId);
}

void TDiskFAManager::IndexLayerUnsafe(const TString& layer, const TBlob& document, const ui32 docId) {
    CHECK_WITH_LOG(!ReadOnly);
    TLayers::const_iterator l = Layers.find(layer);
    if (l == Layers.end()) {
        return;
    }

    if (TLightLayersRegistry::Instance()->IsLightLayer(layer)) {
        auto lightLayer = dynamic_cast<TFALightLayer*>(l->second.Get());
        Y_VERIFY(lightLayer != nullptr);
        lightLayer->IndexUnsafe(document, docId);
    } else {
        l->second->Index(document, docId);
    }
}

bool TDiskFAManager::ReadParsedDoc(const TString& layer, NRTYServer::TParsedDoc& pd, ui32 docId, ui32* size, bool) const {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    TLayers::const_iterator l = Layers.find(layer);
    if (l == Layers.end())
        return false;
    return l->second->ReadParsedDoc(pd, docId, size);
}

TBlob TDiskFAManager::ReadRawDoc(const TString& layer, ui32 docId) const {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    auto l = Layers.find(layer);
    if (l == Layers.end()) {
        return {};
    }
    return l->second->ReadRawDoc(docId);
}

bool TDiskFAManager::IsRemoved(ui32 docId) const {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    TLayers::const_iterator layer1 = Layers.find(NRTYServer::NFullArchive::BaseLayer);
    if (layer1 != Layers.end()) {
        return layer1->second->GetArchive()->IsRemoved(docId);
    }
    for (const auto& layer2 : Layers) {
        if (!layer2.second->GetArchive()->IsRemoved(docId)) {
            return false;
        }
    }
    return true;
}

bool TDiskFAManager::SearchDocument(ui32 docId, const TCgiParameters& cgi, NMetaProtocol::TDocument& doc, const TFAExtraSearchParams* extraParams) const {
    auto parsedDoc = GetDocFast(docId);
    if (!parsedDoc.Defined()) {
        return false;
    }
    CopyDocumentToReport(docId, &parsedDoc->GetDocument(), cgi, doc, extraParams);
    return true;
}

bool TDiskFAManager::SearchDocumentCandidate(TDocIdCandidate docId, const TCgiParameters& cgi, NMetaProtocol::TDocument& doc, const TFAExtraSearchParams* extraParams) const {
    auto parsedDoc = GetDocFast(docId.GetDocId());
    if (!parsedDoc.Defined()) {
        return false;
    }
    if (!docId.IsVerified() && parsedDoc->GetDocument().GetUrl() != docId.GetUrl()) {
        return false;
    }
    CopyDocumentToReport(docId.GetDocId(), &parsedDoc->GetDocument(), cgi, doc, extraParams);
    return true;
}

TParsedDocument::TPtr TDiskFAManager::GetDoc(ui32 docId, NRTYServer::TMessage::TMessageType command, bool) const {
    DEBUG_LOG << "LayersCount: " << Layers.size() << Endl;
    DEBUG_LOG << "Patching doc " << docId << " on restore..." << Endl;
    TParsedDocument::TPtr result;
    NRTYServer::TParsedDoc pd;
    {
        TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
        if (!!FullLayer) {
            if (FullLayer->ReadParsedDoc(pd, docId, nullptr))
                result = FullLayer->DeserializeParsedDocument(pd, command);
        } else {
            for (const auto& layer : Layers) {
                if (!layer.second->ReadParsedDoc(pd, docId, nullptr))
                    continue;
                if (!result) {
                    result = layer.second->DeserializeParsedDocument(pd, command);
                } else {
                    TParsedDocument::TPtr patch = layer.second->DeserializeParsedDocument(pd, NRTYServer::TMessage::DEPRECATED__UPDATE_DOCUMENT);
                    if (patch->GetDocSearchInfo().UrlInitialized() && result->GetDocSearchInfo().UrlInitialized())
                        AssertCorrectIndex(patch->GetDocSearchInfo().GetHash() == result->GetDocSearchInfo().GetHash(), "parts of document is not syncron");
                    result->ApplyPatch(*patch);
                }
            }
        }
    }
    DEBUG_LOG << "Patching doc " << docId << " on restore OK" << Endl;
    return result;
}

TMaybe<NRTYServer::TParsedDoc> TDiskFAManager::GetDocFast(ui32 docId) const {
    CHECK_WITH_LOG(FullLayer);
    TMaybe<NRTYServer::TParsedDoc> parsedDocProto;
    parsedDocProto.ConstructInPlace();
    if (!FullLayer->ReadParsedDoc(*parsedDocProto, docId, nullptr)) {
        parsedDocProto.Clear();
    }
    return parsedDocProto;
}

namespace {
    class TFADocSearchInfoIterator: public NRTYServer::IDocSearchInfoIterator {
    public:
        TFADocSearchInfoIterator(TDiskFAManager::TIterator::TPtr i, const TString& index)
            : Iterator(i)
            , IndexDirectory(index)
        {
            CHECK_WITH_LOG(Iterator);
        }

        TDocSearchInfo Get() const override {
            const NRTYServer::TParsedDoc& pd = Iterator->GetParsedDoc();
            const NRTYServer::TMessage::TDocument& document = pd.GetDocument();
            const ui32 docId = Iterator->GetDocId();
            return TDocSearchInfo(document.GetUrl(), docId, document.GetKeyPrefix(), IndexDirectory);
        }
        bool IsDeleted() const override {
            return false;
        }
        bool IsValid() const override {
            return Iterator->IsValid();
        }
        void Next() override {
            Iterator->Next();
        }

    private:
        const TDiskFAManager::TIterator::TPtr Iterator;
        const TString IndexDirectory;
    };
}

THolder<NRTYServer::IDocSearchInfoIterator> TDiskFAManager::GetDocSearchInfoIterator() const {
    auto p = Layers.find(NRTYServer::NFullArchive::BaseLayer);
    if (p == Layers.end()) {
        WARNING_LOG << "FullArchive base layer not found, trying to use full layer" << Endl;
        p = Layers.find(NRTYServer::NFullArchive::FullLayer);
    }
    if (p == Layers.end()) {
        FAIL_LOG("Cannot find layer with DocSearchInfos");
    }
    return MakeHolder<TFADocSearchInfoIterator>(CreateIterator(p->first), Directory.Basename());
}

void TDiskFAManager::Remap(const TVector<ui32>* remapTable) {
    VERIFY_WITH_LOG(remapTable, "Invalid usage of remapTable");
    CHECK_WITH_LOG(!ReadOnly);
    INFO_LOG << "Remap FULL_ARCHIVE " << Directory.GetPath() << " ..." << Endl;

    VERIFY_WITH_LOG(RemapMutex.Defined(), "RemapMutex does not exist");
    TWriteGuardBase<TMaybe<TRWMutex>> RemapGuard(RemapMutex);
    VERIFY_WITH_LOG(!RemapTable, "remap twice: %s", Directory.GetPath().data());
    size_t newDocsCount = 0;
    for (const auto& i : *remapTable)
        if (i != REMAP_NOWHERE)
            ++newDocsCount;
    RemapTable.Reset(new TRemapTable(*remapTable));
    {
        TIndexMetadataProcessor metaData(Directory);
        metaData->MutableFullArcHeader()->SetStage(NRTYServer::TIndexMetadata::TFullArchiveHeader::REMAP_POSITIONS);
        metaData->MutableFullArcHeader()->SetDocsCount(newDocsCount);
    }
    for (auto& layer : Layers)
        layer.second->Remap(*remapTable);
    TIndexMetadataProcessor(Directory)->MutableFullArcHeader()->SetStage(NRTYServer::TIndexMetadata::TFullArchiveHeader::STARTED);

    TGuard<TMutex> defferedRemovesGuard(DeferredRemovesMutex);
    RemoveDocidsUnsafe(DeferredRemoves);
    DeferredRemoves.clear();
    INFO_LOG << "Remap FULL_ARCHIVE " << Directory.GetPath() << " ...Ok" << Endl;
}

bool TDiskFAManager::Check(const TFsPath& dir, const TSet<TString>& layers, const TRTYFullArchiveConfig& config) {
    const NRTYServer::TIndexMetadata::TFullArchiveHeader header = TIndexMetadataProcessor(dir)->GetFullArcHeader();
    if (header.GetVersion() != FULL_ARC_VERSION && header.GetVersion() != 0) {
        WARNING_LOG << "Invalid version of FULL_ARC in " << dir.GetPath() << ": " << header.GetVersion() << " != " << FULL_ARC_VERSION << Endl;
        return false;
    }
    if (header.GetStage() != NRTYServer::TIndexMetadata::TFullArchiveHeader::FINISHED) {
        WARNING_LOG << "Invalid stage of FULL_ARC in " << dir.GetPath() << ": "
                    << NRTYServer::TIndexMetadata::TFullArchiveHeader::TStage_Name(header.GetStage()) << " != "
                    << NRTYServer::TIndexMetadata::TFullArchiveHeader::TStage_Name(NRTYServer::TIndexMetadata::TFullArchiveHeader::FINISHED) << Endl;
        return false;
    }

    for (auto& layer : layers) {
        if (!TRTYFullArchiveLayer::CheckLayerRight(dir, layer, config.GetLayer(layer))) {
            WARNING_LOG << "FULL_ARC layer checking for " << dir.GetPath() << " failed: "
                        << layer << " broken" << Endl;
            return false;
        }
    }
    return true;
}

bool TDiskFAManager::Repair() {
    TRY {
        const NRTYServer::TIndexMetadata::TFullArchiveHeader header = TIndexMetadataProcessor(Directory)->GetFullArcHeader();
        AssertCorrectIndex(header.GetVersion() == FULL_ARC_VERSION, "cannot repair full archive");

        switch (header.GetStage()) {
            case NRTYServer::TIndexMetadata::TFullArchiveHeader::STARTED:
            case NRTYServer::TIndexMetadata::TFullArchiveHeader::REMAP_POSITIONS:
            case NRTYServer::TIndexMetadata::TFullArchiveHeader::FINISHED:
                RepairByData();
                break;
            default:
                AbortFromCorruptedIndex("invalid full archive state %s", Directory.GetPath().data());
        }

        TIndexMetadataProcessor metaData(Directory);
        metaData->MutableFullArcHeader()->SetVersion(FULL_ARC_VERSION);
        metaData->MutableFullArcHeader()->SetStage(NRTYServer::TIndexMetadata::TFullArchiveHeader::FINISHED);
        return true;
    } CATCH("while repair " + Directory.GetPath())
    return false;
}

void TDiskFAManager::RepairByData() {
    const TRTYFullArchiveConfig& config = *Config.ComponentsConfig.Get<TRTYFullArchiveConfig>(FULL_ARCHIVE_COMPONENT_NAME);

    INFO_LOG << "RepairByData " << Directory.GetPath() << "..." << Endl;
    for (auto& layer : Layers) {
        layer.second->RepairByData(config.GetLayer(layer.first));
    }
    INFO_LOG << "RepairByData " << Directory.GetPath() << "...OK" << Endl;
}

bool TDiskFAManager::DoOpen() {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    VERIFY_WITH_LOG(!IsOpened(), "try to reopen FULLARC manager %s", Directory.GetPath().data());
    TSet<TString> activeLayers;
    for (auto& layer : Layers) {
        activeLayers.insert(layer.first);
    }

    const TRTYFullArchiveConfig& config = *Config.ComponentsConfig.Get<TRTYFullArchiveConfig>(FULL_ARCHIVE_COMPONENT_NAME);
    if (!DocsCount && !Check(Directory, activeLayers, config)) {
        ERROR_LOG << "FULL_ARCHIVE corrupted " << Directory.GetPath() << Endl;
        return false;
    }
    ui32 origDocCount = DocsCount;
    DocsCount = Max<ui32>();
    for (auto& layer : Layers) {
        layer.second->Open(origDocCount);
        ui32 dc = layer.second->GetArchive()->GetDocsCount(true);
        if (DocsCount == Max<ui32>()) {
            DocsCount = dc;
        } else {
            AssertCorrectIndex(DocsCount == dc, "unconsistance layer %s in %s, %i != %i", layer.first.data(), Directory.GetPath().data(), dc, DocsCount);
        }
    }
    if (!ReadOnly) {
        TIndexMetadataProcessor metaData(Directory);
        metaData->MutableFullArcHeader()->SetStage(NRTYServer::TIndexMetadata::TFullArchiveHeader::STARTED);
        metaData->MutableFullArcHeader()->SetDocsCount(DocsCount);
        metaData->MutableFullArcHeader()->SetVersion(FULL_ARC_VERSION);
        metaData->MutableFullArcHeader()->ClearWritedLayers();
        for (const auto& layer : Layers)
            metaData->MutableFullArcHeader()->AddWritedLayers(layer.first);
    }
    return TBaseGeneratorManager::DoOpen();
}

bool TDiskFAManager::DoClose() {
    TReadGuardBase<TMaybe<TRWMutex>> g(RemapMutex);
    if (!IsOpened()) {
        DEBUG_LOG << "Try to close not opened FULLARC manager" << Directory.GetPath() << Endl;
        return false;
    }
    for (auto& layer : Layers) {
        layer.second->Close();
    }
    DocsCount = 0;
    if (!ReadOnly) {
        try {
            TIndexMetadataProcessor metaData(Directory);
            metaData->MutableFullArcHeader()->SetStage(NRTYServer::TIndexMetadata::TFullArchiveHeader::FINISHED);
        } catch (...) {
            const TString info = CurrentExceptionMessage();
            FAIL_LOG("Can't write methadata for DISK-FA: %s", info.data());
        }
    }
    return TBaseGeneratorManager::DoClose();
}

TDiskFAManager::~TDiskFAManager() {
    if (IsOpened())
        Close();
}

bool TDiskFAManager::FullArchiveExists(const TFsPath& dir, const TRTYFullArchiveConfig& config) {
    for (const TString& layer : config.GetActiveLayersFinal()) {
        if (!ExistsLayer(dir, layer)) {
            return false;
        }
    }
    return true;
}

bool TDiskFAManager::ExistsLayer(const TFsPath& dir, const TString& layer) {
    if (TArchiveOwner::Exists(dir / TRTYFullArchiveLayer::GetArchiveName(layer))) {
        return true;
    }
    if (TArchiveOwner::Exists(dir / GetLegacyArchiveName(layer))) {
        return true;
    }
    return false;
}

void TDiskFAManager::RemoveLayer(const TFsPath& dir, const TString& layer) {
    TArchiveOwner::Remove(dir / TRTYFullArchiveLayer::GetArchiveName(layer));
}

bool TDiskFAManager::LayerIsEmpty(const TFsPath& dir, const TString& layer) {
    if (!TArchiveOwner::Empty(dir / TRTYFullArchiveLayer::GetArchiveName(layer))) {
        return false;
    }
    if (!TArchiveOwner::Empty(dir / GetLegacyArchiveName(layer))) {
        return false;
    }
    return true;
}

//TDiskFAManager::TIterator
//TODO(yrum): some copypaste, refactor to remove it
TDiskFAManager::TSampler::TSampler(const TRTYFullArchiveLayer& owner)
    : Owner(owner)
    , DocId(Max<ui32>())
{
    VERIFY_WITH_LOG(owner.IsOpen(), "archive layer not opened");
}

bool TDiskFAManager::TSampler::Seek(ui32 docId) {
    ParsedDoc.Clear();
    Document.Drop();

    TBlob&& blob = Owner.GetArchive()->ReadDocument(docId);
    if (blob.Empty()) {
        DocId = Max<ui32>();
        return false;
    }

    DocId = docId;
    Owner.DeserializeParsedDoc(ParsedDoc, blob);
    return true;
}

const NRTYServer::TParsedDoc& TDiskFAManager::TSampler::GetParsedDoc() const {
    return ParsedDoc;
}

const TParsedDocument::TPtr TDiskFAManager::TSampler::GetDocument(NRTYServer::TMessage::TMessageType command) const {
    if (!IsValid()) {
        return nullptr;
    }
    if (!Document) {
        Document = Owner.DeserializeParsedDocument(ParsedDoc, command);
    }
    return Document;
}

//TDiskFAManager::TIterator
TDiskFAManager::TIterator::TIterator(const TRTYFullArchiveLayer& owner)
    : Owner(owner)
{
    VERIFY_WITH_LOG(owner.IsOpen(), "archive layer not opened");
    Slave = owner.GetArchive()->CreateIterator();
}

const NRTYServer::TParsedDoc& TDiskFAManager::TIterator::GetParsedDoc() const {
    ParsedDoc.Clear();
    Owner.DeserializeParsedDoc(ParsedDoc, GetDocBlob());
    return ParsedDoc;
}

const TParsedDocument::TPtr TDiskFAManager::TIterator::GetDocument(NRTYServer::TMessage::TMessageType command) const {
    if (!IsValid()) {
        return nullptr;
    }
    if (!Document) {
        Document = Owner.DeserializeParsedDocument(GetParsedDoc(), command);
    }
    return Document;
}

NRTYArchive::TArchiveInfo TDiskFAManager::GetArchiveInfo() const {
    NRTYArchive::TArchiveInfo info;
    if (!!FullLayer) {
        FullLayer->GetArchive()->GetInfo(info);
    }
    return info;
}
