#include "report_filter.h"

#include <util/generic/algorithm.h>
#include <util/generic/hash_set.h>
#include <util/generic/vector.h>

namespace NRTYServer {
    namespace NTrie {
        TReportFilter::TReportFilter(ISimpleReportBuilder& output, char prefixDelimiter, size_t prefixLevel, THashMap<TString, ui32> suffixMap)
            : Output(output)
            , SuffixMap(std::move(suffixMap))
            , PrefixLevel(prefixLevel)
            , PrefixDelimiter(prefixDelimiter)
        {
        }

        void TReportFilter::AddDocument(NMetaProtocol::TDocument& doc) {
            TDocumentWrapper wrapper(doc, PrefixDelimiter, PrefixLevel, SuffixMap);
            if (wrapper.GetPrefix().empty()) {
                Output.AddDocument(doc);
            } else {
                Documents.emplace_back().Swap(&doc);
                wrapper.SetDocument(Documents.back());
                Wrappers.push_back(wrapper);
            }
        }

        void TReportFilter::AddDocuments(TList<NMetaProtocol::TDocument> docs) {
            for (auto it = docs.begin(), itEnd = docs.end(); it != itEnd;) {
                auto current = it++;
                TDocumentWrapper wrapper(*current, PrefixDelimiter, PrefixLevel, SuffixMap);
                if (!wrapper.GetPrefix().empty()) {
                    Documents.splice(Documents.end(), docs, current);
                    Wrappers.push_back(wrapper);
                }
            }
            Output.AddDocuments(std::move(docs));
        }

        void TReportFilter::Finalize(TStringBuf attrPrefix) {
            FilterDocuments(Wrappers, Output, attrPrefix);
        }

        void FilterDocuments(TDeque<TDocumentWrapper>& docs, ISimpleReportBuilder& output, TStringBuf attrPrefix) {
            if (docs.empty()) {
                return;
            }
            Sort(docs);
            THashSet<TString> props;
            TString lastPrefix;
            for (auto& wrapper : docs) {
                auto& doc = wrapper.GetDocument();
                if (doc.HasArchiveInfo()) {
                    auto& archiveInfo = *doc.MutableArchiveInfo();
                    auto& attrs = *archiveInfo.MutableGtaRelatedAttribute();
                    size_t dest = 0;
                    size_t nAttrs = attrs.size();
                    if (wrapper.GetPrefix() != lastPrefix) {
                        props.clear();
                        lastPrefix = wrapper.GetPrefix();
                    }
                    for (size_t src = 0; src < nAttrs; ++src) {
                        auto& attr = attrs[src];
                        bool keep = !attr.GetKey().StartsWith(attrPrefix) ||
                            props.insert(attr.GetKey()).second;
                        if (keep) {
                            if (src != dest) {
                                attrs.SwapElements(src, dest);
                            }
                            ++dest;
                        }
                    }
                    if (dest < nAttrs) {
                        attrs.Truncate(dest);
                    }
                }
                output.AddDocument(doc);
            }
        }
    }
}

