#pragma once

#include "search_manager.h"

#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/indexer_core/abstract_model.h>
#include <search/request/data/reqdata.h>

#include <util/system/rwlock.h>

class TCustomSearchManager: public IRealtimeSearchManager {
private:
    const NRTYServer::IIndexComponentManager* IndexComponentManager;
    THashMap<TString, ui32> Url2DocId;
    TVector<TString> DocId2Url;
    TSet<ui32> FreeDocIds;
    TRWMutex Mutex;

    void BanDocImpl(const ui32 docId) {
        FreeDocIds.insert(docId);
        VERIFY_WITH_LOG(docId < DocId2Url.size(), "Incorrect docid in memory indexer");
        Y_VERIFY(Url2DocId.erase(DocId2Url[docId]));
        DocId2Url[docId] = TString();
    }

public:
    TCustomSearchManager(const NRTYServer::IIndexComponentManager* manager)
        : IndexComponentManager(manager)
    {
        AssertCorrectConfig(!!IndexComponentManager, "Trying initialize MemorySearch without IndexComponentManager");
    }

    virtual const NGroupingAttrs::TDocsAttrs* GetDocsAttrs() const override {
        return nullptr;
    }

    virtual const IReadContext::TPtr GetKeysAndPositions() const override {
        return nullptr;
    }

    virtual void BanDoc(const ui32 docId) override {
        TWriteGuard wg(Mutex);
        BanDocImpl(docId);
    }

    virtual bool HasDoc(const TString& url, ui32& docId) const override {
        TReadGuard rg(Mutex);
        const auto docIter = Url2DocId.find(url);
        if (docIter != Url2DocId.end()) {
            docId = docIter->second;
            return !FreeDocIds.contains(docId);
        }
        return false;
    }

    virtual void SignalSearchingStart() override {
    }

    virtual void SearchClose() override {
    }

    virtual void SearchOpen(TAutoPtr<TSearchConfig> /*config*/, const NRTYServer::TIndexerConfig& indexerConfig, TRTYIndexData* /*indexData*/) override {
        for (ui32 i = 0; i < indexerConfig.MaxDocuments; ++i) {
            FreeDocIds.insert(i);
        }
        DocId2Url.resize(indexerConfig.MaxDocuments);
    }

    virtual const IArchiveData* GetArchiveData() const override {
        return nullptr;
    }

    virtual TBaseCommonSearch* GetCommonSearch() override {
        return nullptr;
    }

    virtual ERequestType RequestType(const TCgiParameters& /*cgi*/) const override {
        return RT_Search;
    }

    virtual NMemorySearch::EConsumeDataResult Consume(NRealTime::TIndexedDoc* /*indexedDoc*/, const TDocSearchInfo& docInfo, TCondVar& cvWaitMutex, TMutex& mutexIndex, TQueriesMetrics& /*consumeMetrics*/) override {
        ui32 docId;
        size_t waitCount = 0;
        while (true) {
            for (; FreeDocIds.empty(); ++waitCount) {
                if (waitCount % 100 == 0) {
                    DEBUG_LOG << "Memory indexer consume waiting (waited " << waitCount << " times)" << Endl;
                }
                if (waitCount % 10000 == 0) {
                    WARNING_LOG << "Memory indexer consume waiting (waited " << waitCount << " times)" << Endl;
                }
                cvWaitMutex.WaitT(mutexIndex, TDuration::MilliSeconds(1));
            }

            TWriteGuard wg(Mutex);
            TString urlHash = docInfo.GetMSHash();
            auto it = Url2DocId.find(urlHash);
            if (it != Url2DocId.end()) {
                BanDocImpl(it->second);
            }
            if (!FreeDocIds.empty()) {
                docId = *FreeDocIds.begin();
                FreeDocIds.erase(FreeDocIds.begin());
                Url2DocId[urlHash] = docId;
                VERIFY_WITH_LOG(docId < DocId2Url.size(), "Incorrect docid in memory indexer");
                DocId2Url[docId] = urlHash;
                return NMemorySearch::CDR_OK;
            }
        }
    }
};
