#include <saas/searchproxy/proxy_meta/rearrange/abstract/rearrange.h>

#include <library/cpp/logger/global/global.h>
#include <search/idl/meta.pb.h>
#include <search/session/compression/compression.h>


class TCompressedDataRearrange : public ICustomRearrange {
public:
    TCompressedDataRearrange(): ICustomRearrange()
    {}

    TString GetName() const override {
        return "Compression";
    }

    void DoProcessReplies(TVector<NProxyMeta::TSourceReply>& replies, ICustomReportBuilder&, IReplyContext&) const override {
        for (auto& answer: replies) {
            auto& report = *answer.Report;
            if (report.GroupingSize() == 0)
                continue;
            for (ui32 gr = 0; gr < report.GetGrouping(0).GroupSize() ; ++gr) {
                NMetaProtocol::TGroup* group = report.MutableGrouping(0)->MutableGroup(gr);
                for (ui32 doc = 0; doc < group->DocumentSize(); ++doc) {
                    NMetaProtocol::TArchiveInfo* archiveInfo = group->MutableDocument(doc)->MutableArchiveInfo();
                    for (ui32 i = 0; i < archiveInfo->GtaRelatedAttributeSize(); ++i) {
                        NMetaProtocol::TPairBytesBytes* attr = archiveInfo->MutableGtaRelatedAttribute(i);
                        TVector<TString> keyParts;
                        StringSplitter(attr->GetKey()).Split('_').SkipEmpty().Limit(2).Collect(&keyParts);
                        NMetaProtocol::ECompressionMethod compression;
                        if (keyParts.size() > 1 && TryFromString(keyParts[0], compression)) {
                            TString data = attr->GetValue();
                            attr->SetValue(NMetaProtocol::GetDefaultCompressor().Decompress(data, compression));
                            attr->SetKey(keyParts[1]);
                        }
                    }
                }
            }
        }
    }
};

class TCompressedDataRearrangeFactory : public ICustomRearrangeFactory {
public:
    TCompressedDataRearrangeFactory(const TCustomRearrangeParams&) {
    }

    ICustomRearrange::TPtr CreateRerrangeInstance() const override {
        return new TCompressedDataRearrange();
    }

    TString GetName() const override {
        return "Compression";
    }

private:
    static ICustomRearrangeFactory::TFactory::TRegistrator<TCompressedDataRearrangeFactory> Registrator;
};

ICustomRearrangeFactory::TFactory::TRegistrator<TCompressedDataRearrangeFactory> TCompressedDataRearrangeFactory::Registrator("Compression");
