#include "common.h"

#include <saas/library/searchserver/replier.h>
#include <search/request/data/reqdata.h>
#include <search/session/comsearch.h>
#include <search/output_context/stream_output_context.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/string_utils/quote/quote.h>

#include <util/memory/blob.h>
#include <util/stream/file.h>
#include <util/stream/null.h>
#include <util/string/strip.h>

namespace {
    class TPrefetchContext : public IReplyContext {
    private:
        const TString Request;

        TBlob Buf;
        TNullOutput NullOutput;
        NSearch::TStreamOutputContext NullContext;
        TSearchRequestData RD;
        TInstant RequestStartTime;

    public:
        TPrefetchContext(const TString& request)
            : Request(request)
            , NullContext(NullOutput)
            , RD(Request.data())
            , RequestStartTime(Now())
        {
        }

        void MakeSimpleReply(const TBuffer& /*buf*/, int /*code*/) override {
        }

        bool IsHttp() const override {
            return false;
        }
        NSearch::IOutputContext& Output() override {
            return NullContext;
        }
        TInstant GetRequestStartTime() const override {
            return RequestStartTime;
        }
        TBlob& GetBuf() override {
            return Buf;
        }
        const TSearchRequestData& GetRequestData() const override {
            return RD;
        }
        TSearchRequestData& MutableRequestData() override {
            return RD;
        }
        void AddReplyInfo(const TString&, const TString&) override {}
    };

    class TPrefetchReplier : public TCommonSearchReplier {
    public:
        using TCommonSearchReplier::TCommonSearchReplier;

    protected:
        TDuration GetDefaultTimeout() const override {
            return TDuration::Max();
        }

        void OnRequestExpired(const int /*httpCode*/) override {
            ERROR_LOG << "A 'prefetch' request timed out, took " << (Context->GetRequestStartTime() - Now()) << Endl;
        }

        const TSearchHandlers* GetSearchHandlers() const override {
            return nullptr;
        }
    };

    void PrefetchCommonOneRequest(const TCommonSearch* search, const TString& request, TDuration& elapsed) {
        DEBUG_LOG << "Executing request " << request << " started" << Endl;
        IReplyContext::TPtr context = MakeAtomicShared<TPrefetchContext>(request);
        THolder<ISearchReplier> replier = MakeHolder<TPrefetchReplier>(context, search, nullptr);
        replier.Release()->Reply();

        elapsed = Now() - context->GetRequestStartTime();
        DEBUG_LOG << "Executing request " << request << " took " << elapsed << Endl;
    }

    constexpr TStringBuf prefetchOpts = "&ms=proto&pron=prefetch&nocache=da&skiploadlog=1";

    constexpr TStringBuf externalQueriesOpts = "&timeout=250000&waitall=da";

    TString CreateRequest(const TString& text) {
        return TString(prefetchOpts) + "&text=" + CGIEscapeRet(text);
    }

    const TString DefaultPrefetchRequests[] = {
        "url:*",
        "yandex",
        "google",
        "bing"
    };

    const ui32 DefaultPrefetchCycles = 3;

    constexpr TDuration MaxPrefetchTime = TDuration::Seconds(4);

    constexpr ui32 HeavyRequestsLimit = 24;
}

class NRTYServer::TPrefetcher::TImpl {
private:
    TVector<TString> Queries;

public:
    static TStringBuf TrimQuery(TStringBuf src) {
        src = StripStringLeft(src);
        if (src.StartsWith("# ")) {
            // ignore commented line
            return TStringBuf();
        }
        if (src.StartsWith("http")) {
            src = src.After(':');
            for (int i = 0; i < 3; i++) {
                src = src.After('/');
            }
        }
        if (!src.StartsWith('&')) {
            src = src.After('?');
        }
        return StripStringRight(src);
    }

    void AddRequests(TStringBuf line) {
        // we try to activate various Basesearch features here via "backdoors"
        // There is no guarantee that this stuff keeps working :(
        TString queryFastRank = TString(line) + prefetchOpts + externalQueriesOpts + "&haha=da&pron=onlyfastrank";
        Queries.emplace_back(std::move(queryFastRank));

        TString queryL3 = TString(line) + prefetchOpts + externalQueriesOpts + "&haha=da&pron=nofastrank";
        Queries.emplace_back(std::move(queryL3));

        TString querySnip = TString(line) + prefetchOpts + externalQueriesOpts + "&haha=no";
        Queries.emplace_back(std::move(querySnip));
    }

    bool LoadPlainQueries(const TFsPath& path) {
        TFileInput fi(path);
        for (TString line; fi.ReadLine(line); ) {
            TStringBuf trimmed = TrimQuery(line);
            if (trimmed.size() <= 1) {
                continue;
            }
            AddRequests(trimmed);
        }

        return !Queries.empty();
    }

public:
    TDuration Shoot(const TCommonSearch* search) const {
        ui32 totalQueries = 0;
        TDuration elapsed, totalTime;

        // repeat the shooting plan until a stop condition is met
        bool stopFlag = false;
        do {
            for (const TString& query: Queries) {
                PrefetchCommonOneRequest(search, query, elapsed);
                totalTime += elapsed;
                if ((stopFlag = ++totalQueries >= HeavyRequestsLimit || totalTime >= MaxPrefetchTime))
                    break;
            }
        } while (!stopFlag);
        return totalTime;
    }
};

NRTYServer::TPrefetcher::TPrefetcher(const TFsPath& plainQueriesFile) {
    if (plainQueriesFile.IsDefined() && plainQueriesFile.Exists() && !plainQueriesFile.IsDirectory()) {
        Impl = MakeHolder<TImpl>();

        const bool loaded = Impl->LoadPlainQueries(plainQueriesFile);
        if (!loaded) {
            ERROR_LOG << "Cannot load warmup queries from " << plainQueriesFile << ": no queries loaded" << Endl;
            Impl.Destroy();
        }
    } else {
        if (plainQueriesFile.IsDefined()) {
            ERROR_LOG << "Cannot load warmup queries from " << plainQueriesFile << ": no such file" << Endl;
        }
    }
}

NRTYServer::TPrefetcher::~TPrefetcher() {
}

bool NRTYServer::TPrefetcher::IsEnabled() const {
    return !!Impl;
}

void NRTYServer::TPrefetcher::Prefetch(const TCommonSearch* search) {
    CHECK_WITH_LOG(Impl);
    DEBUG_LOG << "Started warming up CommonSearch" << Endl;
    TDuration totalTime = Impl->Shoot(search);
    if (totalTime > TDuration::Seconds(1)) {
        INFO_LOG << "Warming up CommonSearch took " << totalTime << Endl;
    }
    DEBUG_LOG << "Finished warming up CommonSearch" << Endl;

}

void NRTYServer::PrefetchCommon(const TCommonSearch* search) {
    CHECK_WITH_LOG(search);
    DEBUG_LOG << "Started prefetching CommonSearch" << Endl;
    TDuration elapsed, totalTime;
    for (size_t i = 0; i < DefaultPrefetchCycles && totalTime < MaxPrefetchTime; ++i) {
        for (auto&& text : DefaultPrefetchRequests) {
            PrefetchCommonOneRequest(search, CreateRequest(text), elapsed);
            totalTime += elapsed;
        }
    }
    DEBUG_LOG << "Finished prefetching CommonSearch" << Endl;
}

void NRTYServer::PrefetchCommon(TPrefetcher* prefetcher, const TCommonSearch* search) {
    CHECK_WITH_LOG(search);
    if (!prefetcher || !prefetcher->IsEnabled()) {
        PrefetchCommon(search);
    } else {
        prefetcher->Prefetch(search);
    }
}

