#pragma once

#include <saas/rtyserver/indexer_core/abstract_model.h>
#include <saas/rtyserver/factors/factors_blocks.h>

#include <kernel/indexann/interface/writer.h>

#include <util/system/mutex.h>

class TAnnComponentConfig;
class TSentenceLensProcessor;
class TRTYSentenceLengthWriter;
struct TDirConf;

namespace NIndexAnn {
    class TWriter;
    struct TWriterConfig;
    struct TSentenceParams;
}

namespace NRTYAnn {
    class TStreamDataWriter;
}

class TAnnIndexBuilder: public NRTYServer::IIndexComponentBuilder {
private:
    class TAnnWriter {
        const TAnnComponentConfig& ComponentConfig;
        const TString Prefix;
        TArchiveOwner::TPtr DocumentsStorage;
        TVector<TAtomicSharedPtr<NIndexAnn::TWriter>> AnnWriters;

    public:
        TAnnWriter(const TDirConf& dirConfig,
                   const TAnnComponentConfig& componentConfig,
                   const TRTYServerConfig& globalConfig,
                   const NIndexAnn::TWriterConfig& annWriterConfig,
                   const TString& streamsIndex);

        void Close(const NRTYServer::TBuilderCloseContext& context);

        void StartDoc(ui32 threadId, ui32 docId, ui64 keyPrefix);
        void AddData(ui32 threadId, ui32 region, ui32 streamId, const TString& data);
        void StartSentence(ui32 threadId, const TUtf16String& sentence, const NIndexAnn::TSentenceParams& sentP);
        void FinishDoc(ui32 threadId, ui32 docId);

    private:
        TAtomicSharedPtr<NIndexAnn::IDocDataWriter> CreateAnnDataWriter(
            const NRTYFactors::TAnnFactorsList& streamsInfo
        ) const;
    };

private:
    const TAnnComponentConfig* Config;
    bool UseLanguageFromConf;
    THolder<NIndexAnn::TWriterConfig> AnnConfig;
    THashMap<TString, THolder<TAnnWriter> > Writers;

public:
    TAnnIndexBuilder(const TDirConf& dirConfig, const TRTYServerConfig& rtyConfig, const TString& componentName);
    ~TAnnIndexBuilder();

    NRTYServer::IIndexComponentManager* GetManager() override {
        return nullptr;
    }

    virtual void Index(int threadID, const TParsedDocument& document, const ui32 docId) override;
    virtual bool DoClose(const NRTYServer::TBuilderCloseContext& context) override;
    virtual void InitInteractions(const NRTYServer::IIndexBuildersStorage& /*storage*/) override {}

    virtual bool Start() override {
        return true;
    }

    virtual bool Stop() override {
        return true;
    }
};
