#include "const.h"
#include "config.h"
#include "manager.h"

#include <saas/rtyserver/components/ann/arraydata/iterator.h>
#include <saas/rtyserver/components/ann/storage/iterator.h>
#include <saas/rtyserver/components/ann/storage/accessor.h>
#include <saas/rtyserver/factors/factors_config.h>
#include <saas/rtyserver/config/common_indexers_config.h>
#include <saas/rtyserver/config/indexer_config.h>

#include <kernel/sent_lens/sent_lens.h>

TAnnIndexManager::TAnnIndexManager(const NRTYServer::TManagerConstructionContext& context, const TString& componentName)
    : NRTYServer::IIndexComponentManager(componentName)
    , Config(context.Config.Common.Owner.ComponentsConfig.Get<TAnnComponentConfig>(componentName))
    , Factors(*context.Config.Common.Owner.GetSearcherConfig().Factors)
    , Directory(context.Dir.PathName())
    , IsReadOnly(context.IsReadOnly)
{
    CHECK_WITH_LOG(Config);
    for (const auto& streams : Factors.GetAnnStreams())
        DataStorages[streams.first] = nullptr;
}

TAnnIndexManager::~TAnnIndexManager() {
}

ui32 TAnnIndexManager::GetDocumentsCount() const {
    return Max<ui32>();
}

ui32 TAnnIndexManager::RemoveDocids(const TVector<ui32>& docids) {
    ui32 removed = 0;
    for (auto docId : docids) {
        ui32 count = 0;
        for (auto& dataStorage : DataStorages) {
            if (dataStorage.second)
                count += dataStorage.second->RemoveDocument(docId);
        }
        removed += Min<ui32>(1, count);
    }

    return removed;
}

bool TAnnIndexManager::DoOpen() {
    for (auto& dataStorage : DataStorages) {
        const auto dataAnn = Directory / ("index" + dataStorage.first);
        if (!TArchiveOwner::Check(dataAnn)) {
            ERROR_LOG << "Can't find correct index" << dataStorage.first << " data archive" << Endl;
            return false;
        } else {
            dataStorage.second = TArchiveOwner::Create(
                    dataAnn,
                    Config->MultipartConfig,
                    /*docsCount=*/0u,
                    /*readonly*/IsReadOnly
            );
        }
    }

    return true;
}

bool TAnnIndexManager::DoClose() {
    for (auto& dataStorage : DataStorages) {
        dataStorage.second.Drop();
    }
    return true;
}

TAtomicSharedPtr<NRTYAnn::TAnnDataAccessor> TAnnIndexManager::GetDataBlockAccessor(const TString& type) const {
    return new NRTYAnn::TAnnDataAccessor(*Factors.GetAnnStreams(type));
}

void TAnnIndexManager::InitInteractions(const NRTYServer::IIndexManagersStorage&) {
}

bool TAnnIndexManager::GetDocInfo(const ui32 docId, const TString& type, NJson::TJsonValue& result) const {
    TRTYAnnReader reader(this, type);
    auto iter = reader.CreateIterator(NIndexAnn::THitMask(docId));
    NJson::TJsonValue indexData;
    for (; iter->Valid(); iter->Next()) {
        NJson::TJsonValue& streamData = indexData.AppendValue(NJson::JSON_MAP);
        const auto& hit = iter->Current();
        streamData["break"] = hit.Break();
        streamData["region"] = hit.Region();
        streamData["category"] = 0;
        TString stream = ToString(hit.Stream());
        for (const auto& factor : Factors.GetAnnStreams(type)->GetFactorsList()) {
            if (factor.IndexGlobal == hit.Stream()) {
                stream = factor.Name;
                break;
            }
        }
        streamData["stream"] = stream;
        streamData["data"] = hit.Value();
    }
    result[type] = indexData;
    return true;
}

bool TAnnIndexManager::GetDocInfo(const ui32 docId, NJson::TJsonValue& result) const {
    bool status = false;
    for (const auto& dataStorage : DataStorages) {
        status |= GetDocInfo(docId, dataStorage.first, result);
    }
    return status;
}

bool TAnnIndexManager::HasType(const TString& type) const {
    return DataStorages.contains(type);
}

ui32 TAnnIndexManager::GetFactorAnnHitsLimit() const {
    return GetConfig().FactorAnnHitsLimit;
}

TArchiveOwner::TPtr TAnnIndexManager::GetDataStorage(const TString& type) const {
    auto i = DataStorages.find(type);
    return i == DataStorages.end() ? nullptr : i->second;
}

TRTYAnnReader::TRTYAnnReader(const TAnnIndexManager* manager, const TString& streamIndex)
    : Manager(manager)
    , StreamIndex(streamIndex)
{}

THolder<NIndexAnn::IDocDataIterator> TRTYAnnReader::DoCreateIterator() const {
    const bool remapStreamId = Manager->GetConfig().MapStreamIdToIndex;
    switch (Manager->GetConfig().DataType) {
    case NRTYAnn::EDataFormat::COMPTRIE: {
        TAtomicSharedPtr<NRTYAnn::TAnnDataAccessor> accessor(Manager->GetDataBlockAccessor(StreamIndex));
        return MakeHolder<NRTYAnn::TStreamIterator>(this, accessor, remapStreamId ? accessor.Get() : nullptr);
    }
    case NRTYAnn::EDataFormat::PLAIN_ARRAY:
        return NewArrayDocDataIterator(this, remapStreamId ? this : nullptr);
    }
    Y_VERIFY(false);
}

TBlob TRTYAnnReader::Get(ui32 doc) const {
    return Manager->GetDataStorage(StreamIndex)->GetDocument(doc);
}

bool TRTYAnnReader::HasDoc(ui32 doc) const {
    return !Get(doc).Empty();
}

ui32 TRTYAnnReader::GetGlobalStreamIndex(ui32 streamId) const {
    return Manager->GetFactors().GetAnnStreams(StreamIndex)->GetFactorsInfo()->GetFactor(streamId).IndexGlobal;
}

