#include "manager.h"
#include "component.h"
#include "index.h"

#include <search/base/yx_search.h>

#include <saas/library/mapping/mapping.h>

#include <saas/rtyserver/common/common_rty.h>
#include <saas/rtyserver/components/search/archive/data.h>
#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/common_indexers_config.h>
#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/config/searcher_config.h>
#include <saas/rtyserver/indexer_core/index_component_storage.h>
#include <saas/rtyserver/search/base/search.h>
#include <saas/rtyserver/search/external_search/rty_index_data.h>
#include <saas/rtyserver/search/prefetch/common.h>
#include <saas/rtyserver/search/config/search_config_builder_base.h>
#include <saas/rtyserver/search/factory/factory.h>
#include <saas/util/hex.h>

#include <kernel/index_mapping/index_mapping.h>

#include <library/cpp/logger/global/global.h>

NRTYServer::EUrlIdSource TBaseGeneratorManager::GetUrlIdSource() const {
    return GetIndexComponent()->GetUrlIdSource();
}

void TBaseGeneratorManager::LinkIndexData(IRemapperUrlDocId& /*remapper*/, TRTYIndexData* /*indexData*/) const {
    ++LinkedCounter;
}

void TBaseGeneratorManager::UnlinkIndexData() const {
    --LinkedCounter;
    CHECK_WITH_LOG(LinkedCounter >= 0);
}

const TBaseIndexComponent* TBaseGeneratorManager::GetIndexComponent() const {
    auto result = dynamic_cast<const TBaseIndexComponent*>(TIndexComponentsStorage::Instance().GetComponent(Config.IndexGenerator));
    CHECK_WITH_LOG(result);
    return result;
}

ui32 TBaseGeneratorManager::GetDocumentsCount() const {
    if (CachedDocNum == Max<ui32>() || CachedDocNum == 0) {
        TGuard<TMutex> g(Mutex);
        if (CachedDocNum == Max<ui32>() || CachedDocNum == 0) {
            CachedDocNum = GetDocumentsCountImpl();
        }
    }
    return CachedDocNum;
}

ui32 TBaseGeneratorManager::GetDeletedDocumentsCount() const {
    if (CachedRemovedDocNum == Max<ui32>()) {
        TGuard<TMutex> g(Mutex);
        if (CachedRemovedDocNum == Max<ui32>()) {
            CachedRemovedDocNum = 0;
            for (ui32 i = 0; i < GetDocumentsCount(); ++i) {
                if (IsRemoved(i))
                    ++CachedRemovedDocNum;
            }
        }
    }
    return CachedRemovedDocNum;
}

bool TBaseGeneratorManager::GetDocumentsByKeyPrefix(ui64 kps, TVector<ui32>& docIds) const {
    auto iterator = GetDocSearchInfoIterator();
    if (!iterator) {
        return false;
    }

    for (; iterator->IsValid(); iterator->Next()) {
        const TDocSearchInfo& docSearchInfo = iterator->Get();
        if (docSearchInfo.GetKeyPrefix() == kps && !iterator->IsDeleted()) {
            docIds.push_back(docSearchInfo.GetDocId());
        }
    }
    return true;
}

ui32 TBaseGeneratorManager::RemoveDocids(const TVector<ui32>& docids) {
    ui32 result = DoRemoveDocids(docids);
    {
        TGuard<TMutex> g(Mutex);
        if (CachedRemovedDocNum != Max<ui32>())
            CachedRemovedDocNum += result;
    }
    return result;
}

ui32 TBaseGeneratorManager::RemoveDocidsUnsafe(const TVector<ui32>& docids) {
    ui32 result = DoRemoveDocidsUnsafe(docids);
    {
        TGuard<TMutex> g(Mutex);
        if (CachedRemovedDocNum != Max<ui32>())
            CachedRemovedDocNum += result;
    }
    return result;
}

bool TBaseGeneratorManager::UpdateDoc(ui32 /*docId*/, const TParsedDocument::TPtr /*doc*/) {
    if (Config.IndexGenerator == ComponentName && !Config.IsSecondaryMetaServiceComponent)
        FAIL_LOG("Update not implemented");
    return true;
}

TBaseIndexManager::TBaseIndexManager(const NRTYServer::TManagerConstructionContext& context, const TString& componentName)
    : TBaseGeneratorManager(context.Config.Common.Owner, componentName)
    , Index(context.Index)
    , UseGlobalMapping(context.UseGlobalFileMapping)
    , IsReadOnly(context.IsReadOnly)
    , IndexDir(context.Dir)
    , CheckFrqBordersForRemoveFlag(true)
    , CachedRemovedFrqDocNum(Max<ui32>())
{
}

static void PrecheckBasesearchConfig(const TSearchConfig& base) {
    // check models existance to avoid REFRESH-89 situation
    TFsPath mxNetFile(base.ProtoCollection_.GetMXNetFile());
    AssertCorrectConfig(!mxNetFile.IsDefined() || mxNetFile.Exists(), "File not found: %s", mxNetFile.c_str());
}

void TBaseIndexManager::LinkIndexData(IRemapperUrlDocId& remapper, TRTYIndexData* indexData) const {
    TBaseGeneratorManager::LinkIndexData(remapper, indexData);

    bool createSearch = indexData->GetIndexController()->GetType() != IIndexController::EIndexType::PREPARED; // REFRESH-273
    if (createSearch) {

        if (Y_UNLIKELY(!Config.GetSearcherConfig().CustomBaseSearch.empty())) {
            THolder<NRTYServer::TBaseSearchCreator> BaseSearchCreator(NRTYServer::TBaseSearchCreator::TFactory::Construct(Config.GetSearcherConfig().CustomBaseSearch));
            CHECK_WITH_LOG(BaseSearchCreator.Get() != nullptr);
            BaseSearch = BaseSearchCreator->Create(NRTYServer::TBaseSearchCreator::TParams());
        } else {
            BaseSearch = NRTYServer::CreateBaseSearch(indexData->GetIndexController());
        }

        GetIndexComponent()->TuneSearch(BaseSearch.Get(), remapper, indexData);
        TAutoPtr<TSearchConfig> searchConfig(TSearchConfigBuilderBase(Config).GetSearchConfig(IndexDir, Config.GetSearcherConfig().Enabled));
        PrecheckBasesearchConfig(*searchConfig);
        TInstant baseSearchStart(Now());

        try {
            // changed from AssertCorrectIndex to VERIFY_WITH_LOG because it's better to try to restart and crash
            // than to put FailedIndex to status and do nothing
            VERIFY_WITH_LOG(BaseSearch->SearchOpen(searchConfig, nullptr) == YXOK, "BaseSearch can't start %s", IndexDir.PathName().data());
            AssertCorrectIndex(BaseSearch->BaseIndexData, "BaseSearch not initialized: %s", IndexDir.PathName().data());
        } catch (...) {
            ERROR_LOG << "Error in BaseSearch.SearchOpen() " << CurrentExceptionMessage() << Endl;
            AbortFromCorruptedIndex("Error in BaseSearch::SearchOpen(). BaseSearch can't start %s", IndexDir.PathName().data());
        }

        INFO_LOG << "FinalIndex " << IndexDir.BaseName() << " basesearch start time: " << Now() - baseSearchStart << Endl;
        indexData->SetDocsAttrs(BaseSearch->BaseIndexData->GetDocsAttrs());
    } else {
        CHECK_WITH_LOG(indexData->GetIndexController()->IsSearchable() == false);
    }

    RTY_MEM_LOG("Build IndexData for " + IndexDir.PathName() + " finished");
}

TBaseIndexManager::~TBaseIndexManager() {
}

bool TBaseIndexManager::DoOpen() {
    const TString& index = IndexDir.PathName() + "/index";
    if (!FileIndexFrq) {
        const TString& frq = index + "frq";
        if (NFs::Exists(frq)) {
            FileIndexFrq.Reset(NRTYServer::GetFileMapping(UseGlobalMapping, frq, IsReadOnly));
            FileIndexFrq->Map(0, FileIndexFrq->GetFile().GetLength());
            bool frqLengthIsEven = FileIndexFrq->Length() % 2 == 0;
            CHECK_WITH_LOG(frqLengthIsEven);
        }
    }

    CachedSentCount.Clear();
    CachedWordCount.Clear();

    if (NGroupingAttrs::TDocsAttrs::IndexExists(index) && Config.IndexGenerator == ComponentName) {
        const TString& indexaa = NGroupingAttrs::TDocsAttrs::GetIndexFileName(index);

        const TMemoryMap* global = nullptr;
        if (Index.GetGroupsAttrs()) {
            if (UseGlobalMapping && (global = GetMappedIndexFile(indexaa)))
                Index.GetGroupsAttrs()->InitCommonMode(false, *global);
            else
                Index.GetGroupsAttrs()->InitCommonMode(false, index.data(), IsReadOnly);
        }
    }

    if (Config.IsPrefixedIndex) {
        if (!KpsInfo.Load(IndexDir.PathName()))
            WARNING_LOG << "Cannot load KeyPrefixInfo for " << IndexDir.BaseName() << Endl;
    }

    return TBaseGeneratorManager::DoOpen();
}

bool TBaseIndexManager::DoClose() {
    if (FileIndexFrq) {
        FileIndexFrq.Reset(nullptr);
        if (UseGlobalMapping) {
            NRTYServer::ReleaseFileMapping(IndexDir.PathName() + "/index" "frq");
        }
    }
    return TBaseGeneratorManager::DoClose();
}

ui32 TBaseIndexManager::GetDocumentsCountImpl() const {
    if (!Index.GetYndex()) {
        return GetDocumentsCountFrq();
    } else {
        if (Config.IsPrefixedIndex) {
            return GetYndexDocCount(*Index.GetYndex());
        } else {
            return GetYndexDocCountByUrl(*Index.GetYndex(), false);
        }
    }
}

namespace {
    template <class T>
    void GetDocumentsByKpsImpl(const IKeysAndPositions* kp, ui64 kps, TVector<ui32>& docIds) {
        TString keyUrl;
        if (kps)
            keyUrl = TString::Join('#', Ui64ToHex(kps, 16), "?url=\"");
        else
            keyUrl = "#url=\"";
        TRequestContext rc;
        i32 pos = kp->LowerBound(keyUrl.data(), rc);
        while (true) {
            const YxRecord* rec = EntryByNumber(kp, rc, pos);
            if (!rec)
                break;
            if (!!strncmp(rec->TextPointer, keyUrl.data(), keyUrl.size()))
                break;

            T it;
            it.Init(*kp, rec->Offset, rec->Length, rec->Counter, RH_DEFAULT);

            for (; it.Valid(); ++it) {
                docIds.push_back(it.Doc());
            }

            ++pos;
        }
    }
}

bool TBaseIndexManager::GetDocumentsByKeyPrefix(ui64 kps, TVector<ui32>& docIds) const {
    if (const IKeysAndPositions* yndex = GetYndex()) {
        GetDocumentsByKpsImpl<TPosIterator<>>(yndex, kps, docIds);
        return true;
    } else {
        return TBaseGeneratorManager::GetDocumentsByKeyPrefix(kps, docIds);
    }
}

bool TBaseIndexManager::GetIndexInfo(NJson::TJsonValue& result) const {
    if (!CachedSentCount.Defined() || !CachedWordCount.Defined()) {
        ui64 sentCount = 0;
        ui64 wordCount = 0;

        TSentenceLengthsReader reader(TFsPath(IndexDir.PathName()) / "indexsent");
        for (size_t i = 0; i < reader.GetSize(); i++) {
            if (!IsRemoved(i)) {
                TSentenceLengths lens;
                reader.Get(i, &lens);
                sentCount += lens.size();
                for (size_t j = 0; j < lens.size(); j++) {
                    wordCount += (ui64)lens[j];
                }
            }
        }

        CachedSentCount = sentCount;
        CachedWordCount = wordCount;
    }

    result = NJson::TJsonValue(NJson::JSON_MAP);
    result["SentCount"] = CachedSentCount.GetRef();
    result["WordCount"] = CachedWordCount.GetRef();

    if (Config.IsPrefixedIndex) {
        result["KPS"] = KpsInfo.GenerateJson();
    }

    return true;
}

ui32 TBaseIndexManager::GetDocumentsCountFrq() const {
    CHECK_WITH_LOG(!!FileIndexFrq);
    bool frqLengthIsEven = FileIndexFrq->Length() % 2 == 0;
    CHECK_WITH_LOG(frqLengthIsEven);
    return FileIndexFrq->Length() / 2;
}

ui32 TBaseIndexManager::GetSearchableDocumentsCount() const {
    return GetDocumentsCountFrq() - GetDeletedDocumentsCountFrq();
}

ui32 TBaseIndexManager::GetState() const {
    return BaseSearch ? BaseSearch->BaseIndexData->GetBaseTimestamp() : TBaseGeneratorManager::GetState();
}

bool TBaseIndexManager::IsRemoved(ui32 docid) const {
    if (docid < GetDocumentsCountFrq()) {
        return !!FileIndexFrq && ((i16*)FileIndexFrq->Ptr())[docid] == -1;
    }
    ERROR_LOG << "status=incorrect_docid;docid=" << docid << ";index=" << IndexDir.BaseName() << ";count_docs=" << GetDocumentsCountFrq() << Endl;
    return false;
}

i64 TBaseIndexManager::GroupAttrValue(ui32 docid, const char* attrname, i64 defaultValue) const {
    const NGroupingAttrs::TDocsAttrs* attrs = Index.GetGroupsAttrs();
    if (attrs) {
        ui32 attrnum = attrs->Config().AttrNum(attrname);
        if (attrnum != NGroupingAttrs::TConfig::NotFound) {
            TCategSeries categs;
            attrs->DocCategs(docid, attrnum, categs);
            if (!categs.Empty())
                return categs.GetCateg(0);
        }
    }
    return defaultValue;
}

ui32 TBaseIndexManager::GetDeletedDocumentsCountFrq() const {
    if (CachedRemovedFrqDocNum == Max<ui32>()) {
        TGuard<TMutex> g(MutexFrqCounter);
        if (CachedRemovedFrqDocNum == Max<ui32>()) {
            CachedRemovedFrqDocNum = 0;
            for (ui32 i = 0; i < GetDocumentsCountFrq(); ++i) {
                if (TBaseIndexManager::IsRemoved(i))
                    ++CachedRemovedFrqDocNum;
            }
        }
    }
    return CachedRemovedFrqDocNum;
}

ui32 TBaseIndexManager::DoRemoveDocids(const TVector<ui32>& docids) {
    CachedSentCount.Clear();
    CachedWordCount.Clear();

    ui32 result = 0;
    i16* indexfrq = (i16*)FileIndexFrq->Ptr();
    ui32 length = GetDocumentsCountFrq();
    ui32 countDocs = GetDocumentsCount();
    for (auto i : docids) {
        DEBUG_LOG << "action=delete;docid=" << i << ";index=" << IndexDir.BaseName() << ";" << Endl;
        CHECK_WITH_LOG(i < countDocs);
        CHECK_WITH_LOG((i < length) || (!CheckFrqBordersForRemoveFlag));
        if (i < length) {
            if (indexfrq[i] != -1) {
                ++result;
                indexfrq[i] = -1;
            }
        }
    }

    {
        TGuard<TMutex> g(MutexFrqCounter);
        if (CachedRemovedFrqDocNum != Max<ui32>())
            CachedRemovedFrqDocNum += result;
    }

    return result;
}

namespace {
    i64 GetKeyPrefix(const IArchiveData& archive, ui32 docId) {
        TBlob extInfo = archive.GetDocExtInfo(docId)->UncompressBlob();
        TDocDescr descr;
        descr.UseBlob(extInfo.Data(), (unsigned int)extInfo.Size());
        TDocInfos docInfos;
        descr.ConfigureDocInfos(docInfos);
        TDocInfos::const_iterator i = docInfos.find("prefix");
        return i != docInfos.end() ? HexToui64(i->second) : 0;
    }

    class TIndexDocSearchInfoIterator : public NRTYServer::IDocSearchInfoIterator {
    public:
        TIndexDocSearchInfoIterator(const TBaseIndexManager* manager, const IArchiveData* ad, const TString& index)
            : Manager(manager)
            , AD(ad)
            , IndexDirectory(index)
            , DocId(0)
        {
            CHECK_WITH_LOG(Manager);
            CHECK_WITH_LOG(AD);
        }

        TIndexDocSearchInfoIterator(const TBaseIndexManager* manager, THolder<TArchiveManager>&& ad, const TString& index)
            : TIndexDocSearchInfoIterator(manager, ad.Get(), index)
        {
            OwnedAD = std::move(ad);
        }

        TDocSearchInfo Get() const override {
            Y_ASSERT(IsValid());
            const TString& url = AD->GetUrl(DocId);
            const i64 keyprefix = GetKeyPrefix(*AD, DocId);
            return TDocSearchInfo(url, DocId, keyprefix, IndexDirectory);
        }
        bool IsDeleted() const override {
            Y_ASSERT(IsValid());
            return Manager->IsRemoved(DocId);
        }
        bool IsValid() const override {
            return DocId < Manager->GetDocumentsCount();
        }
        void Next() override {
            Y_ASSERT(IsValid());
            ++DocId;
        }

    private:
        const TBaseIndexManager* Manager;
        const IArchiveData* AD;
        const TString& IndexDirectory;

        THolder<TArchiveManager> OwnedAD;

        ui32 DocId;
    };
}

THolder<NRTYServer::IDocSearchInfoIterator> TBaseIndexManager::GetDocSearchInfoIterator() const {
    if (const IArchiveData* searchArchiveData = GetArchiveData()) {
        return MakeHolder<TIndexDocSearchInfoIterator>(this, searchArchiveData, IndexDir.BaseName());
    } else {
        THolder<TArchiveManager> archiveData = NRTYServer::CreateArchiveData(TFsPath(IndexDir.PathName()) / "index");
        AssertCorrectIndex(!!archiveData, "Cannot create ArchiveData");
        return MakeHolder<TIndexDocSearchInfoIterator>(this, std::move(archiveData), IndexDir.BaseName());
    }
}

bool TBaseIndexManager::GetDocInfo(const ui32 docId, NJson::TJsonValue& result) const {
    result.InsertValue("deleted", IsRemoved(docId));
    result.InsertValue("docid", docId);

    const IArchiveData* ad = GetArchiveData();
    if (ad && ad->IsInArchive(docId)) {
        TDocSearchInfo docSearchInfo(ad->GetUrl(docId), GetKeyPrefix(*ad, docId));
        if (docSearchInfo.UrlInitialized()) {
            result.InsertValue("archive_url", docSearchInfo.GetUrl());
            result.InsertValue("keyprefix", docSearchInfo.GetKeyPrefix());
            result.InsertValue("identifier", docSearchInfo.GetHash().Quote());
        }
    }

    const NGroupingAttrs::TDocsAttrs* attrs = Index.GetGroupsAttrs();
    if (attrs) {
        NJson::TJsonValue groupingAttributes(NJson::JSON_MAP);
        const NGroupingAttrs::TConfig& config = attrs->Config();
        NGroupingAttrs::TDocAttrs docAttributes(config);
        attrs->DocAttrs(docId, &docAttributes);
        for (ui32 i = 0; i < config.AttrCount(); ++i) {
            const TStringBuf name = config.AttrName(i);
            const TCategSeries& values = docAttributes.Attrs(i);
            NJson::TJsonValue valuesJson(NJson::JSON_ARRAY);
            for (auto categ = values.Begin(); categ != values.End(); ++categ)
                valuesJson.AppendValue(*categ);
            groupingAttributes.InsertValue(name, valuesJson);
        }
        result.InsertValue("grouping_attributes", groupingAttributes);
    }
    return true;
}

IKeysAndPositions* TBaseIndexManager::BuildYndex() const {
    auto yndex = NRTYServer::GetYndex(IndexDir.PathName());
    AssertCorrectIndex((!!yndex
            || NRTYServer::HasTextWad(IndexDir.PathName())
            || NRTYServer::HasJupiPrep(IndexDir.PathName())),
            "Cannot build Yndex for %s", IndexDir.BaseName().data());
    return yndex.Release();
}

void TBaseIndexManager::UnlinkIndexData() const {
    if (BaseSearch) {
        BaseSearch->SearchClose();
        BaseSearch.Destroy();
    }
    NRTYServer::ReleaseFileMapping(IndexDir.PathName() + "/indexfrq");
    NRTYServer::ReleaseFileMapping(IndexDir.PathName() + "/indexaa");
    TBaseGeneratorManager::UnlinkIndexData();
}

bool TBaseIndexManager::IsLinked() const {
    CHECK_WITH_LOG((TBaseGeneratorManager::IsLinked() == !!BaseSearch) || !BaseSearch );
    return TBaseGeneratorManager::IsLinked();
}

void TBaseIndexManager::Prefetch() const {
    CHECK_WITH_LOG(IsSearching());
    NRTYServer::PrefetchCommon(BaseSearch.Get());
}

bool TBaseIndexManager::IsSearching() const {
    return !!BaseSearch && BaseSearch->IsSearching();
}

TCommonSearch* TBaseIndexManager::GetSearcher() const {
    return BaseSearch.Get();
}

const IArchiveData* TBaseIndexManager::GetArchiveData() const {
    return BaseSearch ? BaseSearch->GetArchiveData() : nullptr;
}

const IKeysAndPositions* TBaseIndexManager::GetYndex() const {
    return BaseSearch ? &BaseSearch->BaseIndexData->GetTextYR()->YMain() : nullptr;
}
