#pragma once
#include <quality/relev_tools/knn/lib/document_index_reader.h>
#include <quality/relev_tools/knn/lib/hnsw_index_holder.h>
#include <saas/rtyserver/model/manager.h>
#include <kernel/knn_service/component_api.h>
#include <library/cpp/containers/stack_vector/stack_vec.h>
#include <util/generic/ptr.h>
#include <util/generic/hash_set.h>

class IFAManager;
struct TRTYServerConfig;

namespace NHnswGraph {
struct TSearchStats;
}

namespace NRTYServer {
    struct THnswSearchRequest;
    class THnswComponentConfig;

    class THnswIndexManager: public NRTYServer::IIndexComponentManager {
    public:
        THnswIndexManager(TString name, const THnswComponentConfig& config, TAtomicSharedPtr<NKnnOps::THnswIndexHolder> graph, TAtomicSharedPtr<NKnnOps::TDocumentIndexReader> docInex);

        bool DoOpen() override;
        bool DoClose() override;

        ui32 GetDocumentsCount() const override;
        void InitInteractions(const IIndexManagersStorage& storage) override;

        ERTYSearchResult DoSearch(const TRTYSearchRequestContext& context,
                                  ICustomReportBuilder& reportBuilder,
                                  const IIndexController& controller) const override;

        void AddDocToReport(
            size_t requestId,
            const NKnnService::NProtos::TRequest& request,
            ICustomReportBuilder& reportBuilder,
            const TCgiParameters& parameters,
            ui32 docId,
            float dotProductDist) const;

        ui32 DoSearch(const NKnnService::NProtos::TRequest& request,
                      const TCgiParameters& parameters,
                      ICustomReportBuilder& reportBuilder) const;

        struct TSearchResult {
            const THnswIndexManager* Self = nullptr;
            TStackVec<NKnnOps::TNearestResult, 30> Nearest;
            NHnswGraph::TSearchStats HnswStats;
            THashSet<std::pair<TString, TString>> LogLines;
            size_t UsedSearchSize = 0;
            TInstant SearchStartTime;
        };

        void FindNearest(
            const NKnnService::NProtos::TRequest& request,
            TSearchResult& dst) const;

        ui32 RemoveDocids(const TVector<ui32>& docIds) override;

    private:
        TAtomicSharedPtr<NKnnOps::THnswIndexHolder> HnswGraph;
        TAtomicSharedPtr<NKnnOps::TDocumentIndexReader> DocIndex;
        const THnswComponentConfig& Config;
        TString Name;
    };
}
