#include "query_data_trie.h"

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

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

#include <kernel/querydata/cgi/qd_cgi.h>
#include <kernel/querydata/saas/idl/saas_service_options.sc.h>
#include <kernel/querydata/saas/qd_saas_request.h>
#include <kernel/querydata/saas/qd_saas_response.h>

#include <kernel/saas_trie/idl/trie_key.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/protobuf/json/proto2json.h>

#include <library/cpp/string_utils/scan/scan.h>

namespace NProxyMeta {

    const TString NAME = "QueryDataTrie";

    TFakeRequestSender::TFakeRequestSender(NQueryDataSaaS::TSaaSRequestRec::TRef& saasRequestRecRef, TCgiParameters& cgi)
        : SaasRequestRecRef(saasRequestRecRef)
        , Cgi(cgi)
    {
    }

    void TFakeRequestSender::SendRequest(
        const TString& /*client*/,
        const NQueryDataSaaS::TSaasKey& saasKey,
        NQueryDataSaaS::TSaaSRequestRec::TRef requestContext,
        TStringBuf additionalCgiParams)
    {
        SaasRequestRecRef = requestContext;
        Cgi.InsertUnescaped("text", SerializeToCgi(saasKey.Key, false));
        Cgi.ScanAdd(additionalCgiParams);
    }

    NSc::TValue BuildSchemeConfig(const TCustomRearrangeParams& params) {
        NSc::TValue val;
        {
            NQuerySearch::TSaaSServiceOptions<TSchemeTraits> opts{&val};
            opts.MaxKeysInRequest().Set(Max<ui32>());
            opts.MaxRequests().Set(1);
            auto it = params.Parameters.find("DefaultKps");
            if (it != params.Parameters.end()) {
                opts.DefaultKps().Set(FromString<ui64>(it->second));
            } else {
                opts.DefaultKps().Set(0);
            }
            it = params.Parameters.find("SaaSType");
            if (it != params.Parameters.end()) {
                opts.SaaSType().Set(it->second);
            }
            it = params.Parameters.find("OptimizeKeyTypes");
            if (it != params.Parameters.end()) {
                opts.OptimizeKeyTypes().Set(FromString<bool>(it->second));
            }

            for (const auto& kvPair: params.Parameters) {
                if (!kvPair.first.StartsWith("KT_"))
                    continue;
                TStringBuf value = kvPair.second;
                TStringBuf disableBuf;
                TStringBuf lastUniqueBuf;
                if (!value.TrySplit(',', disableBuf, lastUniqueBuf)) {
                    ythrow yexception() << "Cannot parse boolean pair " << kvPair.second;
                }
                auto keyTypeOpt = opts.KeyTypes()[kvPair.first.substr(3)];
                keyTypeOpt.Disable().Set(FromString<bool>(disableBuf));
                keyTypeOpt.LastUnique().Set(FromString<bool>(lastUniqueBuf));

            }
        }
        return val;
    }

    TQueryDataTrieRearrangeConfig::TQueryDataTrieRearrangeConfig(const TCustomRearrangeParams& params) {
        try {
            ServiceOpts.InitFromSchemeThrow(BuildSchemeConfig(params));

            const auto it = params.Parameters.find("MaxDocs");
            if (it != params.Parameters.end()) {
                MaxDocsStr = ToString<ui64>(FromString<ui64>(it->second));
            }
        } catch(...) {
            AbortFromCorruptedConfig("Incorrect config for custom rearrange QueryDataTrie: " + CurrentExceptionMessage());
        }
    }

    TQueryDataTrieRearrange::TQueryDataTrieRearrange(const TQueryDataTrieRearrangeConfig& config)
        : ICustomRearrange()
        , Config(config)
    {
    }

    TString TQueryDataTrieRearrange::GetName() const {
        return NAME;
    }

    bool TQueryDataTrieRearrange::DoFormCgiParameters(TCgiParameters& cgi) {
        NQueryData::TQuerySearchRequestBuilder qdBuilder;
        NQueryData::TRequestRec requestRec;
        qdBuilder.ParseFormedRequest(requestRec, cgi);
        NProxyMeta::TFakeRequestSender fakeSender(SaasRequestRecRef, cgi);
        NQuerySearch::TSaaSRequestStats stats;
        cgi.EraseAll("text");
        GenerateSaasTrieRequests(Default<TString>(), requestRec, Config.ServiceOpts, stats, fakeSender);
        cgi.ReplaceUnescaped("component", "TRIE");
        const TString compSearchStr(TString::Join("comp:TRIE;max_docs:", Config.MaxDocsStr, ";key_type:", ToString(NSaasTrie::ETrieKeyType::ComplexKey)));
        cgi.ReplaceUnescaped("comp_search", compSearchStr);
        cgi.EraseAll("rearr");
        cgi.EraseAll("relev");
        return true;
    }

    void TQueryDataTrieRearrange::DoProcessReplies(TVector<NProxyMeta::TSourceReply>& replies,
        ICustomReportBuilder& builder,
        IReplyContext& context) const
    {
        VERIFY_WITH_LOG(SaasRequestRecRef, "SaasRequestRecRef is empty");
        const NQueryDataSaaS::TSaaSRequestRec& saasRequestRec = *SaasRequestRecRef;
        NQueryData::TQueryData qd;
        NQueryDataSaaS::TSaaSDocument saasDoc;
        NQueryDataSaaS::TProcessingStats stats;

        for (NProxyMeta::TSourceReply& answerSource: replies) {
            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) {
                    saasDoc.Key = group.GetDocument(doc).GetUrl();
                    saasDoc.Records.clear();

                    NMetaProtocol::TArchiveInfo& archiveInfo = *group.MutableDocument(doc)->MutableArchiveInfo();

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

                    for (const auto& attr : archiveInfo.GetGtaRelatedAttribute()) {
                        if (attr.GetKey().StartsWith(NQueryDataSaaS::QD_SAAS_RECORD_KEY_PREFIX)) {
                            saasDoc.Records.emplace_back(attr.GetValue());
                        }
                    }

                    stats += NQueryDataSaaS::ProcessSaaSResponse(qd, saasDoc, saasRequestRec.GetRequestProperties());
                    archiveInfo.Clear();
                }
            }
            builder.ConsumeReport(answer, answerSource.Source);
            answer.ClearGrouping();
        }
        const TCgiParameters& cgi = context.GetCgiParameters();
        const TString& hr = cgi.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());
        }
    }

    TString TQueryDataTrieRearrange::GetConfig() const {
        return Config.ServiceOpts.AsJson();
    }

    class TQueryDataTrieRearrangeFactory : public ICustomRearrangeFactory {
    private:
        TQueryDataTrieRearrangeConfig Config;

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

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

        TString GetName() const override {
            return NAME;
        }

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

    ICustomRearrangeFactory::TFactory::TRegistrator<TQueryDataTrieRearrangeFactory> TQueryDataTrieRearrangeFactory::Registrator(NAME);

} //namespace NProxyMeta
