#include "builder.h"

#include "component.h"
#include "config.h"
#include "manager.h"

#include <kernel/saas_trie/disk_trie_builder.h>
#include <kernel/saas_trie/memory_trie.h>

#include <saas/rtyserver/components/fullarchive/builder.h>

namespace NRTYServer {
    struct TDocSearchInfoAdapter : NSaasTrie::ITrieStorageIterator {
        explicit TDocSearchInfoAdapter(IDocSearchInfoIterator& docIterator)
            : DocIterator(docIterator)
        {
        }

        bool AtEnd() const override {
            return !DocIterator.IsValid();
        }

        TString GetKey() const override {
            return GetTrieKey(DocIterator.Get());
        }

        ui64 GetValue() const override {
            return DocIterator.Get().GetDocId();
        }

        bool Next() override {
            do {
                DocIterator.Next();
            } while (DocIterator.IsValid() && DocIterator.IsDeleted());
            return DocIterator.IsValid();
        }

    private:
        IDocSearchInfoIterator& DocIterator;
    };

    TTrieIndexBuilder::TTrieIndexBuilder(const TString& componentName)
        : IIndexComponentBuilder(componentName)
    {
    }

    bool TTrieIndexBuilder::Start() {
        return true;
    }

    bool TTrieIndexBuilder::Stop() {
        return true;
    }

    TTrieIndexBuilderMemory::TTrieIndexBuilderMemory(const TString& componentName)
        : TTrieIndexBuilder(componentName)
    {
        Storage = NSaasTrie::CreateMemoryTrie();
    }

    void TTrieIndexBuilderMemory::InitInteractions(const IIndexBuildersStorage& /*storage*/) {
    }

    void TTrieIndexBuilderMemory::Index(int /*threadID*/, const TParsedDocument& document, const ui32 docId) {
        auto key = GetTrieKey(document.GetDocSearchInfo());
        CHECK_WITH_LOG(Storage->Put(key, docId)) << "TTrieIndexBuilder: can not index document #" << docId << ", key: " << key << Endl;
    }

    NRTYServer::IIndexComponentManager* TTrieIndexBuilderMemory::GetManager() {
        return new TTrieIndexManager(Storage);
    }

    bool TTrieIndexBuilderMemory::DoClose(const NRTYServer::TBuilderCloseContext& /*context*/) {
        Storage.Reset();
        return true;
    }

    void TTrieIndexBuilderMemory::DoDiscard() {
        Storage.Reset();
    }

    TTrieIndexBuilderDisk::TTrieIndexBuilderDisk(const TString& componentName, TString path)
        : TTrieIndexBuilder(componentName)
        , Path(std::move(path))
    {
    }

    void TTrieIndexBuilderDisk::InitInteractions(const IIndexBuildersStorage& storage) {
        FullArcBuilder = storage.GetBuilder<TRTYFullArchiveBuilder>(FULL_ARCHIVE_COMPONENT_NAME);
        VERIFY_WITH_LOG(FullArcBuilder, "Can not find FULLARC builder");
    }

    void TTrieIndexBuilderDisk::Index(int /*threadID*/, const TParsedDocument& /*document*/, const ui32 /*docId*/) {
    }

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

    bool TTrieIndexBuilderDisk::DoClose(const NRTYServer::TBuilderCloseContext& /*context*/) {
        auto docIterator = FullArcBuilder->GetManager()->GetDocSearchInfoIterator();
        BuildTrieFromDocIterator(Path, Disk, *docIterator);
        return true;
    }

    void TTrieIndexBuilderDisk::DoDiscard() {
        if (NFs::Exists(Path)) {
            NFs::Remove(Path);
        }
    }

    void BuildTrieFromDocIterator(const TString& path, const NSaasTrie::IDiskIO& disk, IDocSearchInfoIterator& docIterator) {
        TDocSearchInfoAdapter iterator(docIterator);
        BuildTrieFromIterator(path, disk, iterator);
    }
}
