#include "client.h"

#include <search/session/compression/report.h>

#include <library/cpp/http/misc/httpcodes.h>
#include <library/cpp/neh/neh.h>
#include <library/cpp/neh/multiclient.h>

#include <util/digest/multi.h>
#include <util/system/spinlock.h>

namespace NSaas {

    class TSearchReply::TSearchAsyncReplyImpl : public TSearchReply::IImpl {
    public:
        TSearchAsyncReplyImpl(NNeh::THandleRef handle, TDuration timeout, bool alwaysUseDataAsResponse=false)
            : Handle(handle)
            , Timeout(timeout)
            , AlwaysUseDataAsResponse(alwaysUseDataAsResponse)
        {
            Y_VERIFY(Handle);
        }

        template<class TResponse>
        TSearchAsyncReplyImpl(const TResponse& response, bool alwaysUseDataAsResponse = false)
            : AlwaysUseDataAsResponse(alwaysUseDataAsResponse)
        {
            FetchUnsafe(response);
        }

        bool IsSucceeded() const override {
            Fetch();
            return Succeded;
        }
        i32 GetCode() const override {
            Fetch();
            return Code;
        }
        const NMetaProtocol::TReport& GetReport() const override {
            Fetch();
            return Report;
        }
        const TString& GetRawReport() const override {
            Fetch();
            return RawReport;
        }
        const TString& GetErrorMessage() const override {
            Fetch();
            return ErrorMessage;
        }

    private:
        void FetchUnsafe(const NNeh::TResponse& response) const {
            try {
                Succeded = !response.IsError();
                Code = response.IsError() ? response.GetErrorCode() : HTTP_OK;
                ErrorMessage = response.GetErrorText();
                RawReport = (response.IsError() && !AlwaysUseDataAsResponse) ? response.GetErrorText() : response.Data;
                if (Succeded && RawReport) {
                    if (!Report.ParseFromString(RawReport)) {
                        Succeded = false;
                    }
                    else {
                        NMetaProtocol::Decompress(Report);
                    }
                }
            } catch (...) {
                Succeded = false;
                Code = HTTP_INTERNAL_SERVER_ERROR;
                RawReport = CurrentExceptionMessage();
            }
            Fetched = true;
        }

        void FetchUnsafe(const NNeh::TResponseRef& response) const {
            if (response) {
                FetchUnsafe(*response);
            }
            Fetched = true;
        }

        void Fetch() const {
            auto guard = Guard(Lock);
            if (Fetched) {
                return;
            }
            NNeh::TResponseRef response = Handle->Wait(Timeout);
            FetchUnsafe(response);
        }

    private:
        const NNeh::THandleRef Handle;
        const TDuration Timeout;
        const bool AlwaysUseDataAsResponse;

        mutable TAdaptiveLock Lock;

        mutable NMetaProtocol::TReport Report;
        mutable TString RawReport;
        mutable i32 Code = -1;
        mutable TString ErrorMessage;
        mutable bool Fetched = false;
        mutable bool Succeded = false;
    };

    class TSearchReply::TSearchNoReplyImpl : public TSearchReply::IImpl {
    public:
        bool IsSucceeded() const override {
            return false;
        }
        i32 GetCode() const override {
            return -1;
        }
        const NMetaProtocol::TReport& GetReport() const override {
            return Default<NMetaProtocol::TReport>();
        }
        const TString& GetRawReport() const override {
            return Default<TString>();
        }
        const TString& GetErrorMessage() const override {
            return Default<TString>();
        }
    };

    class TSearchReply::TSearchErrorReplyImpl : public TSearchReply::IImpl {
    public:
        TSearchErrorReplyImpl(ui32 code, const TString& message)
            : Code(code)
            , Message(message)
        {
            Y_ASSERT(code <= static_cast<ui32>(Max<i32>()));
        }

        bool IsSucceeded() const override {
            return false;
        }
        i32 GetCode() const override {
            return static_cast<i32>(Code);
        }
        const NMetaProtocol::TReport& GetReport() const override {
            static NMetaProtocol::TReport emptyReport;
            return emptyReport;
        }
        const TString& GetRawReport() const override {
            return Message;
        }
        const TString& GetErrorMessage() const override {
            return Message;
        }

    private:
        const ui32 Code;
        const TString Message;
    };

    class TSearchReply::TSearchNoResultsReplyImpl : public TSearchReply::IImpl {
    public:
        TSearchNoResultsReplyImpl(ui32 code)
            : Code(code)
        {
            Y_ASSERT(code <= static_cast<ui32>(Max<i32>()));
        }

        bool IsSucceeded() const override {
            return true;
        }
        i32 GetCode() const override {
            return static_cast<i32>(Code);
        }
        const NMetaProtocol::TReport& GetReport() const override {
            return Default<NMetaProtocol::TReport>();
        }
        const TString& GetRawReport() const override {
            return Default<TString>();
        }
        const TString& GetErrorMessage() const override {
            return Default<TString>();
        }

    private:
        ui32 Code;
    };

    class TSearchReply::TSearchProtoReplyImpl : public TSearchReply::IImpl {
    public:
        TSearchProtoReplyImpl(ui32 code, const char* data, size_t size)
            : Code(code)
        {
            Y_ASSERT(code <= static_cast<ui32>(Max<i32>()));
            Y_ASSERT(data);
            RawReport = TString(data, size);
            Y_PROTOBUF_SUPPRESS_NODISCARD Report.ParseFromString(RawReport);
            NMetaProtocol::Decompress(Report);
        }

        TSearchProtoReplyImpl(ui32 code, const NMetaProtocol::TReport& report)
            : Code(code)
            , RawReport(report.SerializeAsString())
            , Report(report)
        {
        }

        bool IsSucceeded() const override {
            return true;
        }
        i32 GetCode() const override {
            return static_cast<i32>(Code);
        }
        const NMetaProtocol::TReport& GetReport() const override {
            return Report;
        }
        const TString& GetRawReport() const override {
            return RawReport;
        }
        const TString& GetErrorMessage() const override {
            return Default<TString>();
        }

    private:
        ui32 Code;
        TString RawReport;
        NMetaProtocol::TReport Report;
    };

    TSearchReply TSearchReply::FromAsyncHandle(NNeh::THandleRef handle, TDuration timeout, bool alwaysUseDataAsResponse) {
        return TSearchReply{ new TSearchAsyncReplyImpl(handle, timeout, alwaysUseDataAsResponse) };
    }

    TSearchReply TSearchReply::FromAsyncResponse(const NNeh::TResponse& response, bool alwaysUseDataAsResponse /*= false*/) {
        return TSearchReply{ new TSearchAsyncReplyImpl(response, alwaysUseDataAsResponse) };
    }

    TSearchReply TSearchReply::FromAsyncResponse(const NNeh::TResponseRef& response, bool alwaysUseDataAsResponse /*= false*/) {
        return TSearchReply{ new TSearchAsyncReplyImpl(response, alwaysUseDataAsResponse) };
    }

    TSearchReply TSearchReply::FromNoReply()
    {
        return TSearchReply{new TSearchNoReplyImpl()};
    }

    TSearchReply TSearchReply::FromErrorReply(ui32 code, const TString& message)
    {
        return TSearchReply{new TSearchErrorReplyImpl(code, message)};
    }

    TSearchReply TSearchReply::FromNoResultsReply(ui32 code)
    {
        return TSearchReply{new TSearchNoResultsReplyImpl(code)};
    }

    TSearchReply TSearchReply::FromProtoReply(ui32 code, const char* data, size_t size)
    {
        return TSearchReply{new TSearchProtoReplyImpl(code, data, size)};
    }

    TSearchReply TSearchReply::FromProtoReply(ui32 code, const NMetaProtocol::TReport& report)
    {
        return TSearchReply{new TSearchProtoReplyImpl(code, report)};
    }

    const TDuration TSearchClient::DefaultTimeout = TDuration::MilliSeconds(100);

    TSearchReply TSearchClient::SendAsyncQuery(const TQuery& query, TDuration timeout /*= DefaultTimeout*/, bool alwaysUseDataAsResponse /*= false*/) const try {
        const NNeh::THandleRef handle = NNeh::Request(GetFullUrl(query, timeout));
        return handle ? TSearchReply::FromAsyncHandle(handle, timeout, alwaysUseDataAsResponse) : TSearchReply::FromNoReply();
    } catch (...) {
        return TSearchReply::FromErrorReply(HTTP_INTERNAL_SERVER_ERROR, CurrentExceptionMessage());
    }

    TSearchReply TSearchClient::SendQuery(const TQuery& query,
        TDuration timeout,
        const TSelfFlushLogFramePtr& logFrame, const TString& id) const
    {
        TSimpleTimer timer;

        THttpFetcher fetcher(logFrame, "RtySearch", true);
        auto context = CreateSearchContext(fetcher, query, timeout, id);

        if (!context) {
            return TSearchReply::FromNoReply();
        }

        fetcher.Run(THttpFetcher::TAbortOnSuccess());
        return GetSearchReply(context);
    }

    TString TSearchClient::GetFullUrl(const TQuery& query, TDuration timeout) const
    {
        TString urlString;
        TStringOutput urlOut(urlString);
        urlOut << "http://" << Host << ":" << Port << GetFullQueryString(query, timeout);
        return urlString;
    }

    TString TSearchClient::GetFullQueryString(const TQuery& query, const TDuration timeout) const
    {
        TString queryString;
        TStringOutput queryOut(queryString);
        TDuration realTimeout = timeout;
        if (query.GetBalancerTimeout()) {
            realTimeout = Min(timeout, *query.GetBalancerTimeout());
        }
        queryOut << query.BuildQuery()
            << "&service=" << ServiceName
            << "&timeout=" << realTimeout.MicroSeconds()
            << "&ms=proto"
            << "&relev=attr_limit%3D999999999"; // See https://wiki.yandex-team.ru/Users/IvanMorozov/mistakes#attrlimit9999999

        return queryString;
    }

    TSearchClient::TSearchContextPtr TSearchClient::CreateSearchContext(THttpFetcher& fetcher,
        const TQuery& query,
        TDuration timeout,
        const TString& id) const
    {
        TString queryString = GetFullQueryString(query, timeout);
        const size_t urlHash = MultiHash(Host, Port, queryString);

        const int requestId = fetcher.AddRequest(Host.data(), Port, queryString.data(), timeout);
        if (requestId < 0) {
            return nullptr;
        }
        fetcher.SetTaskReqId(requestId, id);

        for (const auto& [key, value] : query.GetHeaders()) {
            fetcher.AddRequestHeader(requestId, key.Data(), value.Data());
        }
        if (TvmManager) {
            fetcher.AddRequestHeader(requestId, NSaas::HEADER_X_YA_SERVICE_TICKET.data(), TvmManager->GetTicket().data());
        }

        return new TSearchContext{fetcher, urlHash, static_cast<size_t>(requestId)};
    }

    TSearchReply TSearchClient::GetSearchReply(const TSearchContextPtr& context) const
    {
        Y_ASSERT(context);
        const IRemoteRequestResult* result = context->Fetcher.GetRequestResult(context->Index);

        if (!result) {
            return TSearchReply::FromNoReply();
        }

        if (HTTP_OK == result->StatusCode()) {
            return TSearchReply::FromProtoReply(result->StatusCode(),
                result->Content().data, result->Content().size);
        } else if (HTTP_NOT_FOUND == result->StatusCode()) {
            return TSearchReply::FromNoResultsReply(result->StatusCode());
        } else {
            const TString reply(result->Content().data, result->Content().size);
            return TSearchReply::FromErrorReply(result->StatusCode(), reply);
        }
    }

    class TAsyncSearchClient::THandleListener : public NRTYServer::TMultiRequesterBase::TRequestData, public NRTYServer::IHandleListener {
    public:
        THandleListener(THolder<ICallback>&& callback, TDuration timeout, bool alwaysUseDataAsResponse)
            : Callback(std::move(callback))
            , Timeout(timeout)
            , AlwaysUseDataAsResponse(alwaysUseDataAsResponse)
        {}

        void OnStart(const NRTYServer::TQueryInfo& /*info*/) const override {
        }

        void OnNotify(const NRTYServer::TQueryInfo& /*info*/, NNeh::TResponseRef& ref) const override {
            if (Callback) {
                Callback->OnReply(TSearchReply::FromAsyncResponse(ref, AlwaysUseDataAsResponse));
            }
        }

        bool OnResend(const NRTYServer::TQueryInfo& /*info*/, TString& /*newAdrr*/, const NNeh::TMessage& /*msg*/, const NNeh::TResponse* resp) const override {
            if (Callback) {
                if (!resp) {
                    return Callback->OnResend(TSearchReply::FromErrorReply(500, "No response"));
                } else {
                    auto code = resp->GetErrorCode();
                    TString msg = resp->Data ? resp->Data : resp->GetErrorText();
                    return Callback->OnResend(TSearchReply::FromErrorReply(code, msg));
                }
            }
            return false;
        }

        bool ReplyIsOk(const NRTYServer::TQueryInfo& info, const NNeh::TResponse& rep) const override {
            return (!rep.IsError() || rep.GetErrorCode() < 500 && rep.GetErrorCode() > 0) && (!Callback || Callback->ReplyIsOk(info, TSearchReply::FromAsyncResponse(rep)));
        }

        void OnCancel(const NRTYServer::TQueryInfo& /*info*/, const TString& reason) const override {
            if (Callback) {
                Callback->OnReply(TSearchReply::FromErrorReply(500, reason));
            }
        }

        TDuration GetRecvTimeout() const override {
            return Timeout;
        }
    private:
        THolder<ICallback> Callback;
        TDuration Timeout;
        bool AlwaysUseDataAsResponse;
    };

    TAsyncSearchClient::TAsyncSearchClient(const TString& serviceName, const TString& host, const ui16 port, ui32 attemps, const TVector<TDuration>& resendDurations, std::shared_ptr<NSaas::TTvmManager> tvmManager)
        : TSearchClient(serviceName, host, port, tvmManager)
        , Requester(MakeHolder<NRTYServer::TMultiRequesterBase>(attemps, resendDurations))
    {}

    void TAsyncSearchClient::Stop() {
        Requester->Stop(true);
    }

    void TAsyncSearchClient::Start(ui32 threads) {
        Requester->Start(threads);
    }

    void TAsyncSearchClient::AddQuery(const TQuery& query, THolder<ICallback>&& callback, TDuration timeout, bool alwaysUseDataAsResponse) const {
        auto listener = MakeIntrusive<THandleListener>(std::move(callback), timeout, alwaysUseDataAsResponse);
        Requester->Send(NNeh::TMessage::FromString(GetFullUrl(query, timeout)), listener.Get(), listener);
    }

}
