#include <saas/tools/querysearch/qd_send_request/params.pb.h>

#include <search/idl/meta.pb.h>
#include <search/session/compression/report.h>

#include <kernel/querydata/idl/querydata_structs_client.pb.h>
#include <kernel/querydata/saas/qd_saas_key.h>
#include <kernel/saas_trie/idl/saas_trie.pb.h>
#include <kernel/saas_trie/idl/trie_key.h>


#include <library/cpp/getoptpb/getoptpb.h>
#include <library/cpp/http/client/client.h>
#include <library/cpp/logger/global/global.h>
#include <library/cpp/protobuf/json/proto2json.h>
#include <library/cpp/string_utils/base64/base64.h>
#include <library/cpp/json/json_writer.h>

#include <google/protobuf/text_format.h>

#include <util/stream/file.h>
#include <util/stream/input.h>
#include <util/stream/output.h>

namespace {
    const TStringBuf CGI_PARAMS = "?comp_search=comp%3ATRIE;max_docs%3A100500;key_type%3Acomplex_key_packed&component=TRIE&meta_search=first_found&ms=proto&normal_kv_report=1&skip-wizard=1&sp_meta_search=multi_proxy&rearr=QueryDataTrie_off&text=";

    TString SendRequest(const TString& url) {
        NHttp::TFetchOptions options;
        options.RetryCount = 2;
        options.RetryDelay = TDuration::Seconds(1);
        options.Timeout = TDuration::Seconds(5);
        auto result = NHttp::Fetch({url, options});
        if (result->Code >= 400 && result->Code < 500 && result->Code != 429) {
            result = NHttp::Fetch({url + "&hr=1", options});
        }
        Y_ENSURE(result->Code == 200, "Request failed, code = " << result->Code << ", response: " << result->Data);
        return result->Data;
    }

    void PrintReport(const TString& reportString, bool json) {
        NMetaProtocol::TReport report;
        Y_VERIFY(report.ParseFromString(reportString));
        NMetaProtocol::Decompress(report);
        NJson::TJsonValue values(NJson::JSON_ARRAY);
        for (auto& grouping : report.GetGrouping()) {
            for (auto& group : grouping.GetGroup()) {
                for (auto& doc : group.GetDocument()) {
                    if (doc.HasArchiveInfo()) {
                        for (auto& attr : doc.GetArchiveInfo().GetGtaRelatedAttribute()) {
                            if (attr.GetKey().StartsWith("QDSaaS:")) {
                                NQueryData::TSourceFactors sf;
                                Y_ENSURE(sf.ParseFromString(Base64Decode(attr.GetValue())), "Cannot parse TSourceFactors protobuf");
                                if (!json) {
                                    Cout << sf.Utf8DebugString() << Endl;
                                } else {
                                    NJson::TJsonValue value;
                                    NProtobufJson::Proto2Json(sf, value);
                                    values.AppendValue(std::move(value));
                                }
                            }
                        }
                    }
                }
            }
        }
        if (json) {
            NJson::WriteJson(&Cout, &values, true);
        }
    }
}

int main(int argc, const char* argv[]) {
    NGetoptPb::TGetoptPbSettings getoptSettings;
    getoptSettings.DumpConfig = false;
    NUtils::TParams params = NGetoptPb::GetoptPbOrAbort(argc, argv, getoptSettings);
    if (params.GetVerbose()) {
        InitGlobalLog2Console(TLOG_DEBUG);
    }
    try {
        NSaasTrie::TComplexKey key;

        if (!params.GetProtoKey().empty()) {
            DEBUG_LOG << "Read the key from a prototext file: " << params.GetProtoKey() << Endl;
            TIFStream input(params.GetProtoKey());
            ::google::protobuf::TextFormat::ParseFromString(input.ReadAll(), &key);

        } else {
            using namespace NQueryDataSaaS;

            DEBUG_LOG << "Build the key from pieces" << Endl;
            TString mainKey;
            bool keyHasRegions = false;

            auto fillRealm = [&key, &mainKey](ESaaSSubkeyType subkeyType, const auto& values) {
                if (values.empty()) {
                    return false;
                }
                auto strType = ToString(static_cast<int>(subkeyType));
                mainKey += "." + strType;
                key.AddKeyRealms(strType);
                auto realm = key.AddAllRealms();
                realm->SetName(strType);
                for (auto& subkeyValue : values) {
                    realm->AddKey("\t" + subkeyValue);
                }
                return true;
            };

            const auto& allSubkeyTypes = GetAllTrieSubkeysForSearch();
            for (auto subkeyType : allSubkeyTypes) {
                switch (subkeyType) {
                    case ESaaSSubkeyType::SST_IN_KEY_QUERY_STRONG:
                        fillRealm(subkeyType, params.GetQueryStrong());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_QUERY_DOPPEL:
                        fillRealm(subkeyType, params.GetQueryDoppel());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_USER_ID:
                        fillRealm(subkeyType, params.GetUserId());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_USER_LOGIN_HASH:
                        fillRealm(subkeyType, params.GetUserLoginHash());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_USER_REGION:
                        keyHasRegions = fillRealm(subkeyType, params.GetUserRegion());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_STRUCT_KEY:
                        fillRealm(subkeyType, params.GetStructKey());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_URL:
                        fillRealm(subkeyType, params.GetUrl());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_OWNER:
                        fillRealm(subkeyType, params.GetOwner());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_URL_MASK:
                        fillRealm(subkeyType, params.GetUrlMask());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_ZDOCID:
                        fillRealm(subkeyType, params.GetDocId());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_QUERY_DOPPEL_TOKEN:
                        fillRealm(subkeyType, params.GetQueryDoppelToken());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_QUERY_DOPPEL_PAIR:
                        fillRealm(subkeyType, params.GetQueryDoppelPair());
                        break;
                    case ESaaSSubkeyType::SST_IN_KEY_URL_NO_CGI:
                        fillRealm(subkeyType, params.GetUrlNoCgi());
                        break;
                    default:;
                }
            }

            if (params.GetKeyType().empty()) {
                key.SetMainKey(mainKey);
            } else {
                DEBUG_LOG << "Override auto-generated MainKey with " << params.GetKeyType() << Endl;
                key.SetMainKey(params.GetKeyType());
            }
            if (keyHasRegions) {
                Y_ENSURE(mainKey.EndsWith(".5"), "UserRegion has to be the last part of the key");
                DEBUG_LOG << "UserRegion requires LastRealmUnique flag" << Endl;
                key.SetLastRealmUnique(true);
            }

            if (params.GetBanfilterMasks()) {
                Y_ENSURE(mainKey.StartsWith(".7"), "banfilter masks work only if Url is the first part of the key");
                DEBUG_LOG << "Enable banfilter-like masks" << Endl;
                key.SetUrlMaskPrefix(".9");
            }
        }

        DEBUG_LOG << "Key:\n" << key.Utf8DebugString() << Endl;

        const auto textKey = NSaasTrie::SerializeToCgi(key, true);
        const auto requestUrl = TString::Join(params.GetHost(), '/', params.GetService(), CGI_PARAMS, textKey, '&', params.GetExtraCgi());

        DEBUG_LOG << "Send a request: " << requestUrl << Endl;
        const auto reportString = SendRequest(requestUrl);
        DEBUG_LOG << "Got an answer." << Endl;

        PrintReport(reportString, params.GetJson());

        DEBUG_LOG << "Finished." << Endl;

        return 0;

    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
        return 1;
    }
}


