#include "builders_storage.h"
#include "builder.h"
#include "manager.h"

#include <saas/rtyserver/components/fullarchive/manager.h>
#include <saas/rtyserver/config/const.h>
#include <library/cpp/logger/global/global.h>
#include <library/cpp/json/json_value.h>


const NRTYServer::IIndexComponentManager* TRTYIndexManagersStorage::GetManagerImplNoLock(const TString& name) const {
    for (size_t i = 0; i < Managers.size(); ++i) {
        const TManagerRec& managerRec = Managers[i];
        if (managerRec.first == name) {
            return managerRec.second.Get();
        }
    }
    return nullptr;
}

bool TRTYIndexManagersStorage::HasIndexManager() const {
    return !!GetManager<TBaseGeneratorManager>(RTYConfig.IndexGenerator);
}

const TBaseGeneratorManager* TRTYIndexManagersStorage::GetIndexManager() const {
    const TBaseGeneratorManager *manager = GetManager<TBaseGeneratorManager>(RTYConfig.IndexGenerator);
    CHECK_WITH_LOG(manager);
    return manager;
}

const IDDKManager* TRTYIndexManagersStorage::GetDDKManager() const {
    const IGetterDDKManager* ddkmanager = GetManager<IGetterDDKManager>(RTYConfig.DDKManager);
    VERIFY_WITH_LOG(ddkmanager, "Invalid usage of DDKManager : %s", RTYConfig.DDKManager.c_str());
    INFO_LOG << "TRTYIndexManagersStorage uses " << RTYConfig.DDKManager << "'s DDKManager " << Endl;
    return &ddkmanager->GetDDKManager();
}

const IUrlToDocIdManager* TRTYIndexManagersStorage::GetUrlToDocIdManager() const {
    const IUrlToDocIdManager* manager = GetManager<IUrlToDocIdManager>(RTYConfig.UrlToDocIdManager);
    DEBUG_LOG << "TRTYIndexManagersStorage uses " << RTYConfig.UrlToDocIdManager << "'s UrlToDocIdManager " << Endl;
    return manager;
}

void TRTYIndexManagersStorage::KillManagers() {
    Close();
    TWriteGuard g(ReadWriteMutex);
    Managers.clear();
}

void TRTYIndexManagersStorage::InitInteractions() {
    for (const auto& [_, manager]: Managers) {
        manager->InitInteractions(*this);
    }
}

const NRTYServer::IIndexComponentManager* TRTYIndexManagersStorage::GetManagerImpl(const TString& name) const {
    TReadGuard g(ReadWriteMutex);
    return GetManagerImplNoLock(name);
}

void TRTYIndexManagersStorage::RegisterManager(const TString& name, THolder<NRTYServer::IIndexComponentManager> manager) {
    if (!manager) {
        return;
    }
    TWriteGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(GetManagerImplNoLock(name) == nullptr, "RegisterManager %s failed: already registered", name.data());
    Managers.emplace_back(name, std::move(manager));
}

void TRTYIndexManagersStorage::CheckDocumentsNumber(ui32 docsCount) const {
    TReadGuard g(ReadWriteMutex);
    for (size_t i = 0; i < Managers.size(); ++i) {
        const TManagerRec& managerRec = Managers[i];
        ui32 count = managerRec.second->GetDocumentsCount();
        if (count != Max<ui32>())
            AssertCorrectIndex(count == docsCount, "For %s component's manager: IIndexComponentManager::CheckDocumentsCount mismatch %d != %d", managerRec.first.data(), count, docsCount);
    }
}

bool TRTYIndexManagersStorage::UpdateDoc(ui32 docId, const TParsedDocument::TPtr doc) {
    TReadGuard g(ReadWriteMutex);
    bool result = false;
    for (TManagerList::iterator i = Managers.begin(); i != Managers.end(); ++i)
        result |= i->second->UpdateDoc(docId, doc);
    return result;
}

ui32 TRTYIndexManagersStorage::RemoveDocids(const TVector<ui32>& docids) {
    TReadGuard g(ReadWriteMutex);
    ui32 result = 0;
    for (TManagerList::reverse_iterator i = Managers.rbegin(); i != Managers.rend(); ++i)
        result = Max<ui32>(result, i->second->RemoveDocids(docids));
    return result;
}

ui32 TRTYIndexManagersStorage::MarkDocIdsForDeleteUnsafe(const TVector<ui32>& docids, ui32 marker) {
    TReadGuard g(ReadWriteMutex);
    ui32 result = 0;
    for (TManagerList::iterator i = Managers.begin(); i != Managers.end(); ++i)
        result = Max<ui32>(result, i->second->MarkDocIdsForDeleteUnsafe(docids, marker));
    return result;
}

bool TRTYIndexManagersStorage::DoOpen() {
    TReadGuard g(ReadWriteMutex);
    bool rc = true;
    for (size_t i = 0; i < Managers.size(); ++i) {
        const TManagerRec& managerRec = Managers[i];
        RTY_MEM_LOG("Manager " + managerRec.first + " open start");
        bool result = managerRec.second->Open();
        if (result) {
            RTY_MEM_LOG("Manager " + managerRec.first + " open finished");
        }
        else {
            WARNING_LOG << managerRec.first.data() << " IIndexComponentManager::Open failed" << Endl;
            rc = false;
        }
    }
    if (rc && HasIndexManager()) {
        // prefetch cached document counters in TBaseGeneratorManager
        auto docCount = GetIndexManager()->GetSearchableDocumentsCount();
        Y_UNUSED(docCount);
    }
    return rc;
}

bool TRTYIndexManagersStorage::DoClose() {
    TReadGuard g(ReadWriteMutex);
    bool rc = true;
    for (size_t i = 0; i < Managers.size(); ++i) {
        const TManagerRec& managerRec = Managers[Managers.size() - i - 1];
        RTY_MEM_LOG("Manager " + managerRec.first + " close start");
        if (Callbacks)
            Callbacks->OnBeforeManagerClose(managerRec.second.Get());
        bool result = managerRec.second->Close();
        if (Callbacks)
            Callbacks->OnAfterManagerClose(managerRec.second.Get());
        if (result) {
            RTY_MEM_LOG("Manager " + managerRec.first + " close finished");
        } else {
            WARNING_LOG << managerRec.first.data() << " IIndexComponentManager::Close failed" << Endl;
            rc = false;
        }
    }
    return rc;
}

ui32 TRTYIndexManagersStorage::GetDocumentsCount() const {
    return GetIndexManager()->GetDocumentsCount();
}

void TRTYIndexManagersStorage::InitInteractions(const NRTYServer::IIndexManagersStorage& /*storage*/) {
    VERIFY_WITH_LOG(false, "UNSUPPORTED TRTYIndexManagersStorage IIndexComponentManager::InitInteractions");
}

bool TRTYIndexManagersStorage::GetDocInfo(const ui32 docId, NJson::TJsonValue& result) const {
    TReadGuard g(ReadWriteMutex);
    bool success = false;
    for (size_t i = 0; i < Managers.size(); ++i) {
        const TManagerRec& managerRec = Managers[i];
        NJson::TJsonValue json;
        if (managerRec.second->GetDocInfo(docId, json)) {
            result.InsertValue(managerRec.first, json);
            success = true;
        }
    }
    return success;
}

void TRTYIndexManagersStorage::Prefetch() const {
    TReadGuard g(ReadWriteMutex);
    for (auto&& m : Managers) {
        m.second->Prefetch();
    }
}

THolder<NRTYServer::IDocSearchInfoIterator> TRTYIndexManagersStorage::GetDocSearchInfoIterator() const {
    TReadGuard g(ReadWriteMutex);
    const NRTYServer::IIndexComponentManager* manager = nullptr;
    if (!manager || !manager->IsOpened()) {
        manager = GetManagerImpl(FULL_ARCHIVE_COMPONENT_NAME);
    }
    if (!manager || !manager->IsOpened()) {
        manager = GetIndexManager();
    }
    if (!manager) {
        ERROR_LOG << "Cannot create DocSearchInfoIterator: no viable managers found" << Endl;
        return nullptr;
    }
    if (!manager->IsOpened()) {
        ERROR_LOG << "Cannot create DocSearchInfoIterator: " << manager->GetComponentName() << " manager is not opened" << Endl;
        return nullptr;
    }
    return manager->GetDocSearchInfoIterator();
}

NRTYServer::IIndexComponentBuilder* TRTYIndexBuildersStorage::GetBuilderImplNoLock(const TString& name) const {
    for (size_t i = 0; i < Builders.size(); ++i) {
        const TBuilderRec& builderRec = Builders[i];
        if (builderRec.first == name) {
            return builderRec.second.Get();
        }
    }
    return nullptr;
}

void TRTYIndexBuildersStorage::KillBuilders() {
    TWriteGuard g(ReadWriteMutex);
    Discard();
    Builders.clear();
    Managers.KillManagers();
}

NRTYServer::IIndexComponentBuilder* TRTYIndexBuildersStorage::GetBuilderImpl(const TString& name) const
{
    TReadGuard g(ReadWriteMutex);
    return GetBuilderImplNoLock(name);
}

void TRTYIndexBuildersStorage::RegisterBuilder(const TString& name, THolder<NRTYServer::IIndexComponentBuilder> builder) {
    if (!builder) {
        return;
    }
    TWriteGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(!IsRunning, "Incorrect RegisterBuilder usage");
    VERIFY_WITH_LOG(GetBuilderImplNoLock(name) == nullptr, "RegisterBuilder %s failed: already registered", name.data());
    Builders.emplace_back(name, std::move(builder));
}

bool TRTYIndexBuildersStorage::HasIndexBuilder() const {
    return GetBuilderImpl(RTYConfig.IndexGenerator);
}

TBaseGeneratorBuilder* TRTYIndexBuildersStorage::GetIndexBuilder() const {
    return GetBuilder<TBaseGeneratorBuilder>(RTYConfig.IndexGenerator);
}

NRTYServer::IIndexComponentManager* TRTYIndexBuildersStorage::GetManager() {
    FAIL_LOG("UNSUPPORTED TRTYIndexBuildersStorage IIndexComponentBuilder::GetManager");
    return nullptr;
}

void TRTYIndexBuildersStorage::Index(int threadID, const TParsedDocument& document, const ui32 docId) {
    TReadGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(IsRunning, "Incorrect TRTYIndexBuildersStorage::Index usage");
    for (size_t i = 0; i < Builders.size(); ++i) {
        const TBuilderRec& builderRec = Builders[i];
        builderRec.second->Index(threadID, document, docId);
    }
}

bool TRTYIndexBuildersStorage::DoClose(const NRTYServer::TBuilderCloseContext& context) {
    TReadGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(!IsRunning, "Incorrect TRTYIndexBuildersStorage::Close usage");
    for (size_t i = 0; i < Builders.size(); ++i) {
        const TBuilderRec& builderRec = Builders[i];
        try {
            if (Callbacks)
                Callbacks->OnBeforeBuilderClose(builderRec.second.Get(), builderRec.first);
            bool result = builderRec.second->Close(context);
            if (Callbacks)
                Callbacks->OnAfterBuilderClose(builderRec.second.Get(), builderRec.first);
            if (!result && !(!!context.RigidStopSignal && *context.RigidStopSignal)) {
                FAIL_LOG("%s IIndexComponentBuilder::Close failed", builderRec.first.data());
            }
        } catch (...) {
            const TString exceptionMessage = CurrentExceptionMessage();
            FAIL_LOG("%s builder closing failed: %s", builderRec.first.data(), exceptionMessage.data());
        }
    }
    return Managers.Close();
}

void TRTYIndexBuildersStorage::DoDiscard() {
    for (auto& builder : Builders) {
        builder.second->Discard();
    }
}

bool TRTYIndexBuildersStorage::Start(bool initManagersInterations) {
    TWriteGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(!IsRunning, "Incorrect TRTYIndexBuildersStorage::Start usage");
    for (const auto& component : ManagersForOpen) {
        for (auto& builder : Builders) {
            if (builder.first == component) {
                Managers.RegisterManager(component, THolder(builder.second->GetManager()));
            }
        }
    }

    if (initManagersInterations) {
        Managers.InitInteractions();
    }

    Managers.Open();
    for (size_t i = 0; i < Builders.size(); ++i) {
        const TBuilderRec& builderRec = Builders[i];
        RTY_MEM_LOG("Builder " + builderRec.first + " starting...");
        VERIFY_WITH_LOG(builderRec.second->Start(), "%s IIndexComponentBuilder::Start failed", builderRec.first.data());
        RTY_MEM_LOG("Builder " + builderRec.first + " starting... OK");
    }
    IsRunning = true;
    return true;
}

bool TRTYIndexBuildersStorage::Stop() {
    TWriteGuard g(ReadWriteMutex);
    VERIFY_WITH_LOG(IsRunning, "Incorrect TRTYIndexBuildersStorage::Stop usage");
    for (size_t i = 0; i < Builders.size(); ++i) {
        const TBuilderRec& builderRec = Builders[i];
        VERIFY_WITH_LOG(builderRec.second->Stop(), "IIndexComponentBuilder::Stop failed for %s", builderRec.first.data());
    }
    IsRunning = false;
    return true;
}

void TRTYIndexBuildersStorage::InitInteractions(const NRTYServer::IIndexBuildersStorage& /*storage*/) {
    VERIFY_WITH_LOG(false, "UNSUPPORTED TRTYIndexBuildersStorage IIndexComponentBuilder::InitInteractions");
}
