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

#include <saas/library/daemon_base/daemon/messages.h>
#include <saas/library/querydata/report.h>

#include <library/cpp/protobuf/json/proto2json.h>
#include <search/idl/meta.pb.h>

struct TQueryDataRearrangeConfig {
    TString Namespace;
    TString SourceName;
    NQueryData::EKeyType KeyType;

    TQueryDataRearrangeConfig(const TCustomRearrangeParams& params) {
        NQueryData::TKeyTypeDescr kt = NQueryData::KeyTypeDescrByName(params.Parameters.at("KeyType"));

        KeyType = (NQueryData::EKeyType)kt.KeyType;
        Namespace = params.Parameters.at("Namespace");
        SourceName = params.Parameters.at("SourceName");

        if (KeyType == NQueryData::KT_COUNT || Namespace.empty() || SourceName.empty()) {
            AbortFromCorruptedConfig("Incorrect config for custom rearrange QueryData");
        }
    }
};


class TQueryDataRearrange : public ICustomRearrange {
public:
    TQueryDataRearrange(const TQueryDataRearrangeConfig& config)
        : ICustomRearrange()
        , Config(config)
    {
    }

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

    void DoProcessReplies(TVector<NProxyMeta::TSourceReply>& replies, ICustomReportBuilder& builder,
        IReplyContext& context) const override
    {
        const TString& specKeyName = context.GetCgiParameters().Get("key_name");

        NSaas::TQueryDataReportBuilder qdBuilder;
        for (ui32 source = 0; source < replies.size(); ++source) {
            NProxyMeta::TSourceReply& answerSource = replies[source];
            NMetaProtocol::TReport& answer = *answerSource.Report;
            if (answer.GroupingSize() == 0)
                continue;

            for (ui32 gr = 0; gr < answer.GetGrouping(0).GroupSize(); ++gr) {
                NMetaProtocol::TGroup& group = *answer.MutableGrouping(0)->MutableGroup(gr);
                for (ui32 doc = 0; doc < group.DocumentSize(); ++doc) {
                    const TString url = group.GetDocument(doc).GetUrl();
                    NMetaProtocol::TArchiveInfo& archiveInfo = *group.MutableDocument(doc)->MutableArchiveInfo();

                    if (archiveInfo.GtaRelatedAttributeSize() == 0) {
                        continue;
                    }

                    NSaas::TQueryDataReportBuilder::TDocData qdDocData =
                        qdBuilder.AddDocumentDataSet(url, Config.Namespace, Config.SourceName, Config.KeyType);

                    for (ui32 i = 0; i < archiveInfo.GtaRelatedAttributeSize(); ++i) {
                        NMetaProtocol::TPairBytesBytes attr = archiveInfo.GetGtaRelatedAttribute(i);

                        if (specKeyName) {
                            if (attr.GetKey().StartsWith("_SK_")) {
                                if (attr.GetKey() == "_SK_" + specKeyName) {
                                    qdDocData.UpdateSourceKey(attr.GetValue());
                                }
                                continue;
                            }
                        }

                        qdDocData.Add(attr.GetKey(), attr.GetValue());
                    }
                    archiveInfo.Clear();
                }
            }
            builder.ConsumeReport(answer, answerSource.Source);
            answer.ClearGrouping();
        }

        const NQueryData::TQueryData& qd = qdBuilder.GetFinalReport();
        const TString& hr = context.GetCgiParameters().Get("hr");
        if (IsTrue(hr) || hr == "json") {
            TStringStream ss;
            NProtobufJson::Proto2Json(qd, ss);
            builder.AddReportProperty("QueryData.debug", ss.Str());
        } else {
            builder.AddReportProperty("QueryData.debug", qd.SerializeAsString());
        }
    }

private:
    const TQueryDataRearrangeConfig& Config;
};

class TQueryDataRearrangeFactory : public ICustomRearrangeFactory {
private:
    TQueryDataRearrangeConfig Config;

public:
    TQueryDataRearrangeFactory(const TCustomRearrangeParams& params)
        : Config(params)
    {
    }

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

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

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

ICustomRearrangeFactory::TFactory::TRegistrator<TQueryDataRearrangeFactory> TQueryDataRearrangeFactory::Registrator("QueryData");

