#include "trie_search_request.h"

#include <kernel/saas_trie/config.h>
#include <kernel/saas_trie/duplicate_key_filter.h>
#include <kernel/saas_trie/idl/saas_trie.pb.h>
#include <kernel/saas_trie/idl/trie_key.h>
#include <kernel/saas_trie/trie_iterator_chain.h>

#include <util/string/cast.h>

using namespace NSaasTrie;

namespace NRTYServer {
    TTrieSearchRequest::TTrieSearchRequest(const TCgiParameters& cgi,
                                           const TMap<TString, TString>* compParams,
                                           bool sortKeys,
                                           const TString& uniqueDimension,
                                           TString propertyPrefix,
                                           ISimpleReportBuilder& reportBuilder)
        : ReportBuilder(reportBuilder)
        , PropertyPrefix(std::move(propertyPrefix))
    {
        ETrieKeyType keyType = ETrieKeyType::Text;
        if (compParams != nullptr) {
            auto it = compParams->find("max_docs");
            if (it != compParams->end()) {
                MaxDocs = FromString<ui32>(it->second);
            }
            it = compParams->find("key_type");
            if (it != compParams->end()) {
                keyType = FromString<ETrieKeyType>(it->second);
            }
        }

        auto range = cgi.Range("text");
        if (keyType == ETrieKeyType::Text) {
            for (const auto& text: range) {
                ui64 kps = cgi.Has("sgkps") ? FromString<ui64>(cgi.Get("sgkps")) : 0;
                Texts.emplace_back(NSaasTrie::GetTrieKey(kps, text));
            }
            IteratorCount = 1;
            UniqueOffset = 1;
        } else {
            size_t nKeys = std::distance(range.begin(), range.end());
            ComplexKeys.reserve(nKeys);
            ComplexKeysUnique.reserve(nKeys);

            for (const auto& text: range) {
                TComplexKey complexKeyProto;
                DeserializeFromCgi(complexKeyProto, text, keyType == ETrieKeyType::ComplexKeyPacked);
                auto key = PreprocessComplexKey(complexKeyProto, sortKeys, uniqueDimension);
                if (key.LastDimensionUnique && !PropertyPrefix.empty()) {
                    ComplexKeysUnique.emplace_back(std::move(key));
                } else {
                    ComplexKeys.emplace_back(std::move(key));
                }
            }
            UniqueOffset = ComplexKeys.empty() ? 0 : 1;
            IteratorCount = UniqueOffset + ComplexKeysUnique.size();
        }
    }

    bool TTrieSearchRequest::IsValid() const {
        return MaxDocs > 0 && (!Texts.empty() || !ComplexKeys.empty() || !ComplexKeysUnique.empty());
    }

    THolder<ITrieStorageIterator> TTrieSearchRequest::CreateIterator(const ITrieStorageReader& storage) const {
        if (!Texts.empty()) {
            if (Texts.size() == 1) {
                return storage.CreatePrefixIterator(Texts.front());
            }
            TTrieIteratorVector iterators;
            iterators.reserve(Texts.size());
            for (auto& text : Texts) {
                iterators.emplace_back(storage.CreatePrefixIterator(text));
            }
            return DecorateTrieIteratorVector(std::move(iterators));
        }
        if (IteratorIndex < IteratorCount) {
            if (IteratorIndex >= UniqueOffset) {
                return FilterDuplicateKeys(CreateTrieComplexKeyIterator(storage, ComplexKeysUnique[IteratorIndex - UniqueOffset]));
            }
            if (ComplexKeys.size() == 1) {
                return FilterDuplicateKeys(CreateTrieComplexKeyIterator(storage, ComplexKeys.front()));
            }
            TTrieIteratorVector iterators;
            iterators.reserve(ComplexKeys.size());
            for (auto& complexKey : ComplexKeys) {
                iterators.emplace_back(CreateTrieComplexKeyIterator(storage, complexKey));
            }
            return FilterDuplicateKeys(DecorateTrieIteratorVector(std::move(iterators)));
        }
        Y_ENSURE(false, "No iterators left");
    }

    ui32 TTrieSearchRequest::GetMaxDocs() const {
        return MaxDocs;
    }

    ui32 TTrieSearchRequest::DecreaseMaxDocs(ui32 foundCount) {
        MaxDocs = MaxDocs < foundCount ? 0 : MaxDocs - foundCount;
        return MaxDocs;
    }

    bool TTrieSearchRequest::NextIterator() {
        if (ReportFilter.Defined()) {
            ReportFilter->Finalize(PropertyPrefix);
            ReportFilter.Clear();
        }
        if (++IteratorIndex >= IteratorCount) {
            return false;
        }
        if (IteratorIndex < UniqueOffset) {
            return true;
        }
        auto& realms = ComplexKeysUnique[IteratorIndex - UniqueOffset].SortedRealms;
        Y_ENSURE(realms.size() > 0);
        auto& lastRealm = realms.back();
        THashMap<TString, ui32> suffixMap;
        for (ui32 i = 0, imax = lastRealm.size(); i < imax; ++i) {
            suffixMap.emplace(lastRealm[i], i);
        }
        ReportFilter.ConstructInPlace(ReportBuilder, '\t', realms.size() - 1, std::move(suffixMap));
        return true;
    }

    ISimpleReportBuilder& TTrieSearchRequest::GetReport() {
        if (ReportFilter.Defined()) {
            return *ReportFilter;
        }
        return ReportBuilder;
    }
}
