#include "broadcast.h"

#include <saas/searchproxy/common/cgi.h>
#include <library/cpp/logger/global/global.h>
#include <search/request/data/reqdata.h>
#include <search/idl/meta.pb.h>

#include <library/cpp/json/json_reader.h>
#include <library/cpp/neh/neh.h>
#include <library/cpp/protobuf/json/proto2json.h>

#include <google/protobuf/text_format.h>

namespace {
    class TBroadcastScheduler : public NSearchMapParser::ISearchMapScannerCallback {
    public:
        TBroadcastScheduler(const TString& service, const TString& request)
            : Service(service)
            , Request(request)
        {
        }

        inline TVector<NNeh::THandleRef>&& GetHandles() {
            return std::move(Handles);
        }
        inline TMap<TString, NSearchMapParser::TSearchMapHost>&& GetHosts() {
            return std::move(Hosts);
        }

        bool OnService(const NSearchMapParser::TSearchMapService& info) override {
            return Service ? info.Name == Service : true;
        }
        void OnHost(const NSearchMapParser::TSearchMapHost& host,
            const NSearchMapParser::TSearchMapReplica& /*replica*/,
            const NSearchMapParser::TSearchMapService& /*service*/
        ) override {
            TString addr = "http://" + host.GetDescription() + "/";
            try {
                NNeh::TMessage message = NNeh::TMessage::FromString("http://" + host.GetDescription() + "/" + Request);
                Handles.push_back(NNeh::Request(message));
                addr = message.Addr;
            } catch(...) {
                WARNING_LOG << "Skipped for broadcast: " << addr << Endl;
            }
            Hosts.insert(std::make_pair(addr, host));
        }

    private:
        const TString Service;
        const TString Request;
        TVector<NNeh::THandleRef> Handles;
        TMap<TString, NSearchMapParser::TSearchMapHost> Hosts;
    };

    inline bool TryParseMetaProtobuf(const TString& data, NJson::TJsonValue& result) {
        google::protobuf::LogSilencer logSilencer;
        NMetaProtocol::TReport report;
        if (!report.ParseFromString(data) && !google::protobuf::TextFormat::ParseFromString(data, &report)) {
            return false;
        }

        NProtobufJson::Proto2Json(report, result);
        return true;
    }
    inline bool TryParseJson(const TString& data, NJson::TJsonValue& result) {
        return NJson::ReadJsonFastTree(data, &result);
    }
    inline bool TryParsePlainData(const TString& data, NJson::TJsonValue& result) {
        result = data;
        return true;
    }

    inline NJson::TJsonValue ResponseToJson(const TString& data) {
        NJson::TJsonValue result;
        bool parsed = false;
        if (!parsed) {
            parsed = TryParseMetaProtobuf(data, result);
        }
        if (!parsed) {
            parsed = TryParseJson(data, result);
        }
        if (!parsed) {
            parsed = TryParsePlainData(data, result);
        }
        return result;
    }
}

NSearchProxy::TBroadcastContextPtr NSearchProxy::TBroadcaster::Broadcast(const TString& service, const TString& request, const TDuration timeout) const {
    NSearchProxy::TBroadcastContextPtr result = MakeIntrusive<TBroadcastContext>(service, request);
    Broadcast(service, request, timeout, [result](const TBackendReply& reply) {
        result->Replies.push_back(reply);
        result->Requested++;
        if (reply.Failed) {
            result->Failed++;
        }
    });
    return result;
}

void NSearchProxy::TBroadcaster::Broadcast(const TString& service, const TString& request, const TDuration timeout, std::function<void(const TBackendReply&)> f) const {
    TBroadcastScheduler scheduler(service, request);
    SearchMap.Scan(scheduler);

    const auto& handles = scheduler.GetHandles();
    const auto& hosts = scheduler.GetHosts();
    const TInstant deadline = Now() + timeout;

    TSet<TString> answered;
    for (auto&& handle : handles) {
        const NNeh::TResponseRef res = handle->Wait(deadline);
        if (!res) {
            continue;
        }

        auto host = hosts.find(res->Request.Addr);
        Y_ASSERT(host != hosts.end());
        NSearchProxy::TBackendReply reply;
        reply.Host = host != hosts.end() ? host->second : NSearchMapParser::TSearchMapHost();
        reply.Failed = res->IsError();
        reply.Data = res->IsError() ? res->GetErrorText() : res->Data;
        f(reply);
        answered.insert(host->first);
    }

    for (auto&& host : hosts) {
        if (answered.contains(host.first)) {
            continue;
        }

        NSearchProxy::TBackendReply reply;
        reply.Host = host.second;
        reply.Failed = true;
        reply.Data = "No response";
        f(reply);
    }
}

NJson::TJsonValue NSearchProxy::ToJson(const TBackendReply& reply) {
    NJson::TJsonValue result;
    result["Host"] = reply.Host.Name;
    result["SearchPort"] = reply.Host.SearchPort;
    result["Shards"] = reply.Host.Shards.ToString();
    if (reply.Failed) {
        result["Error"] = reply.Data;
    } else {
        result["Response"] = ResponseToJson(reply.Data);
    }
    return result;
}
