#include "component.h"

#include "builder.h"
#include "config.h"
#include "manager.h"
#include "searcher.h"
#include "trie_search_request.h"

#include <saas/rtyserver/config/common_indexers_config.h>
#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/config/searcher_config.h>

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

#include <library/cpp/containers/mh_heap/abstract_merger.h>

namespace NRTYServer {
    class TTrieParsedEntity: public TParsedDocument::TParsedEntity {
    public:
        TTrieParsedEntity(TConstructParams& params)
            : TParsedDocument::TParsedEntity(params)
        {
        }
    };

    class TTrieComponentParser: public TComponentParser {
    public:
        void Parse(TParsingContext& /*context*/) const override {
        }
    };

    TTrieComponent::TTrieComponent(const TRTYServerConfig& config)
        : IIndexComponent(config.ComponentsSet.contains(TrieComponentName))
        , Config(config)
        , TrieConfig(*config.ComponentsConfig.Get<TTrieComponentConfig>(TrieComponentName))
    {
        CHECK_WITH_LOG(&TrieConfig);
        IndexFiles.insert(TIndexFile(TrieIndexName, true, TIndexFile::ppDisable));
        MultiSearcher = CreateTrieSearcher(Config.IndexGenerator, TrieConfig.GetSearchThreadCount());
        CHECK_WITH_LOG(MultiSearcher);
    }

    TString TTrieComponent::GetName() const {
        return TrieComponentName;
    }

    bool TTrieComponent::IsCommonSearch() const {
        return false;
    }

    IIndexComponent::TPriorityInfo TTrieComponent::GetPriority() const {
        auto priority = IIndexComponent::GetPriority();
        priority.BuilderPriority = FULLARC_BUILDER_PRIORITY - 1;
        priority.ComponentPriority = FULLARC_COMPONENT_PRIORITY - 1;
        priority.ManagerPriority = FULLARC_MANAGER_PRIORITY - 1;
        return priority;
    }

    THolder<IIndexComponentBuilder> TTrieComponent::CreateBuilder(const TBuilderConstructionContext& context) const {
        if (context.Config.GetType() == "memory") {
            return MakeHolder<TTrieIndexBuilderMemory>(GetName());
        }
        return MakeHolder<TTrieIndexBuilderDisk>(GetName(), GetTrieIndexPath(context.TempDir.PathName()));
    }

    THolder<IIndexComponentManager> TTrieComponent::CreateManager(const TManagerConstructionContext& context) const {
        if (context.IndexType != IIndexController::FINAL) {
            return nullptr;
        }
        auto path = GetTrieIndexPath(context.Dir.PathName());
        auto trie = NSaasTrie::OpenDiskTrie(path, Disk, Config.GetSearcherConfig().ArchivePolicy == ARCHIVE_POLICY_MAPMEMLOCK);
        VERIFY_WITH_LOG(trie, "Can not open trie %s", path.data());
        return MakeHolder<TTrieIndexManager>(std::move(trie));
    }

    IComponentParser::TPtr TTrieComponent::BuildParser() const {
        return new TTrieComponentParser{};
    }

    IParsedEntity::TPtr TTrieComponent::BuildParsedEntity(IParsedEntity::TConstructParams& params) const {
        return new TTrieParsedEntity(params);
    }

    bool TTrieComponent::DoAllRight(const TNormalizerContext& context) const {
        try {
            return CheckDiskTrieIsValid(GetTrieIndexPath(context.Dir.PathName()), Disk);
        } catch (...) {
            ERROR_LOG << "TTrieComponent::DoAllRight failed: " << CurrentExceptionMessage() << Endl;
            return false;
        }
    }

    void TTrieComponent::CheckAndFix(const TNormalizerContext& context) const {
        auto path = GetTrieIndexPath(context.Dir.PathName());
        if (!DoAllRight(context)) {
            auto docIterator = context.Managers.GetDocSearchInfoIterator();
            Y_VERIFY(docIterator);
            BuildTrieFromDocIterator(path, Disk, *docIterator);
        }
    }

    const IIndexComponent::TIndexFiles& TTrieComponent::GetIndexFiles() const {
        return IndexFiles;
    }

    struct TTrieIterable {
        using value_type = TStringBuf;

        TTrieIterable(ui32 sourceId, const TString& indexDirectory, const NSaasTrie::IDiskIO& disk, bool lockMemory)
            : SourceId(sourceId)
        {
            auto path = GetTrieIndexPath(indexDirectory);
            Trie = NSaasTrie::OpenDiskTrie(path, disk, lockMemory);
            VERIFY_WITH_LOG(Trie, "Can not open trie %s", path.data());
            Iterator = Trie->CreateIterator();
            if (!Iterator->AtEnd()) {
                Key = Iterator->GetKey();
            }
        }
        void Restart() {
            // stub for library/container/mh_heap/mh_heap
        }
        bool Valid() const {
            return !Iterator->AtEnd();
        }
        TTrieIterable& operator++() {
            Key = Iterator->Next() ? Iterator->GetKey() : TString{};
            return *this;
        }
        TStringBuf Current() const {
            return {Key.data(), Key.size()};
        }
        ui32 GetDocId() const {
            return static_cast<ui32>(Iterator->GetValue());
        }
        ui32 GetSourceId() const {
            return SourceId;
        }

    private:
        TAtomicSharedPtr<NSaasTrie::ITrieStorageReader> Trie;
        THolder<NSaasTrie::ITrieStorageIterator> Iterator;
        TString Key;
        const ui32 SourceId;
    };

    bool TTrieComponent::DoMerge(const TMergeContext& context) const {
        auto& ctx = context.Context;
        Y_VERIFY(ctx.Decoder);

        TVector<THolder<NSaasTrie::ITrieStorageWriter>> outputs(ctx.Dests.size());
        for (size_t i = 0, imax = outputs.size(); i < imax; ++i) {
            outputs[i] = CreateTrieBuilder(GetTrieIndexPath(ctx.Dests[i]), Disk);
        }

        TAbstractMerger<TTrieIterable> merger;
        bool lockMemory = Config.GetSearcherConfig().ArchivePolicy == ARCHIVE_POLICY_MAPMEMLOCK;
        for (ui32 i = 0, imax = static_cast<ui32>(ctx.Sources.size()); i < imax; ++i) {
            merger.CreateIterable(i, ctx.Sources[i], Disk, lockMemory);
        }
        merger.Restart();

        for (; merger.Valid(); ++merger) {
            auto& input = merger.Current();
            auto addr = ctx.Decoder->Decode(input.GetSourceId(), input.GetDocId());
            if (!addr.IsRemoved()) {
                outputs[addr.ClusterId]->Put(input.Current(), addr.DocId);
            }
        }
        return true;
    }

    void TTrieComponent::SearchCustom(const TVector<IIndexController::TPtr>& controllers,
                                      ICustomReportBuilder& report,
                                      const TRTYSearchRequestContext& context) const {
        TTrieSearchRequest request(context.CgiParams(), context.GetParameters(TrieComponentName),
                                   TrieConfig.NeedToSortComplexKey(), TrieConfig.GetUniqueDimension(),
                                   TrieConfig.GetPropertyPrefix(), report);
        if (!request.IsValid()) {
            return;
        }
        while (request.NextIterator() && request.GetMaxDocs() > 0) {
            MultiSearcher->DoSearch(request, context.CgiParams(), controllers);
        }
    }
}
