#pragma once

#include <saas/rtyserver/factors/factor_calculate.h>
#include <saas/rtyserver/components/l2/l2_disk_manager.h>
#include <saas/rtyserver/search/cgi/rty_search_statistics.h>

namespace NRTYServer {
    TString ComposeSearchKey(const TString& name, const TString& version);

    class TDssmManager: public TL2ComponentDiskManagerBase {
    public:
        using TL2ComponentDiskManagerBase::TL2ComponentDiskManagerBase;

        TString GetDocumentEmbeddingsJson(ui32 docId) const;

        THolder<NRTYFactors::IFactorCalculate> GetDecompressedEmbedding(TCalcFactorsContext& ctx, const TRTYFunctionCtx& uctx, TArrayRef<const NRTYFactors::IFactorCalculate*> args) const;
        THolder<NRTYFactors::IFactorCalculate> GetDecompressedEmbeddings(TCalcFactorsContext& ctx, const TRTYFunctionCtx& uctx, TArrayRef<const NRTYFactors::IFactorCalculate*> args) const;
        THolder<NRTYFactors::IFactorCalculate> GetDecompressedEmbeddingsWithTags(TCalcFactorsContext& ctx, const TRTYFunctionCtx& uctx, TArrayRef<const NRTYFactors::IFactorCalculate*> args) const;
        /**
         * Function to export pairs of (embedding, version) from storage.
         *
         * Main usage is sparse embeddings, that does not exist for big fraction of documents
         * since other functions are throwing exceptions when embedding with given version is not found.
         *
         * Arguments:
         *  1. Model name
         *  2. Compression algorithm (will be used for all versions of embedding)
         *  3-12. Versions to extract (up to 10 versions per model name)
         *
         *  Should be used together with following functions:
         *    - `preload_doc_embeddings` - for saving mapping version => embedding in cache
         *    - `doc_embedding_exists` - to check if embedding with given version exists
         *    - `load_cached_embedding` - to further load cached embedding for computation
         *
         *  Example usage can be found in tests (lookup TestSparseEmbeddings)
         */
        THolder<NRTYFactors::IFactorCalculate> GetDecompressedEmbeddingsByVersion(TCalcFactorsContext& ctx, const TRTYFunctionCtx& uctx, TArrayRef<const NRTYFactors::IFactorCalculate*> args) const;
        void GetExportedFunctions(NRTYFeatures::TImportedFunctionsBuilder& exports) const override;
    private:
        THolder<NRTYFactors::IFactorCalculate> GetDecompressedEmbeddingInternal(
            TCalcFactorsContext& ctx, TRTYSearchStatistics& searchStatistics, TArrayRef<const NRTYFactors::IFactorCalculate*> args, bool checkSingle, bool extractTags) const;
    };
}
