#include "module.h"


#include <balancer/modules/exp_common/exp_common.h>

#include <balancer/kernel/custom_io/chunkio.h>
#include <balancer/kernel/custom_io/queue.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_pire.h>

#include <util/generic/ptr.h>
#include <util/generic/strbuf.h>
#include <util/stream/zlib.h>
#include <util/string/join.h>

using namespace NConfig;
using namespace NSrvKernel;

// TODO(uranix): merge with NExpCommon::ExpHeaders()
static const TFsm HEADER_WHITELIST_REGEXP(JoinSeq("|", {
        "Host",
        "User-Agent",
        "Cookie",
        "X-Req-Id",
        "X-Forwarded-For",
        "X-Forwarded-For-Y",
        "X-Region-City-Id",
        "X-Yandex-ExpSplitParams",
        "X-Yandex-ExpConfigVersion(-Pre)?",
        "X-Yandex-ExpBoxes(-Pre)?",
        "X-Yandex-ExpBoxes-Crypted(-Pre)?",
        "X-Yandex-ICookie",
        "X-Yandex-Internal-Request",
        "X-Yandex-LoginHash",
        "X-Yandex-LogstatUID",
        "X-Yandex-RandomUID",
        "X-Yandex-Suspected-Robot",
        "X-L7-EXP",
        "X-Yandex-Internal-SrcRwr",
        "X-Yandex-TCP-Info"
    }), TFsm::TOptions().SetSurround(false).SetCaseInsensitive(true));

static const TFsm HEADER_ONLY_LENGTH_REGEXP(JoinSeq("|", {
        "X-Yandex-ExpFlags(-Pre)?",
    }), TFsm::TOptions().SetSurround(false).SetCaseInsensitive(true));


Y_TLS(remote_log) {
    TTls(IModule* remoteLogStorage, const TString& noRemoteLogFile, const TString& levelCompressFile, size_t queueLimit,
        IWorkerCtl* process, TSharedCounter inQueue, TSharedCounter dropped)
        : RemoteLoggerProps(remoteLogStorage)
        , RemoteLogger(MakeHolder<NExpCommon::TRemoteLogger>(&RemoteLoggerProps, *process, queueLimit,
            std::move(inQueue), std::move(dropped)))
    {
        RemoteLoggerProps.Init(process, noRemoteLogFile, levelCompressFile);
    }

    NExpCommon::TRemoteLoggerProperties RemoteLoggerProps;
    THolder<NExpCommon::TRemoteLogger> RemoteLogger;
    TInstant NextRequest = Now();
};

MODULE_WITH_TLS_BASE(remote_log, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
        , InQueue_(Control->SharedStatsManager().MakeGauge("remote_log-in_queue").AllowDuplicate().Build())
        , Dropped_(Control->SharedStatsManager().MakeGauge("remote_log-dropped").AllowDuplicate().Build())
    {
        Config->ForEach(this);

        if (Delay_ != TDuration::Zero()) {
            Throttled_ = true;
        }

        if (!Submodule_) {
            ythrow TConfigParseError() << "no mainline module configured";
        }

        if (!RemoteLogStorage_.Get()) {
            ythrow TConfigParseError() << "no log storage module configured";
        }
    }

private:
    START_PARSE {
        ON_KEY("delay", Delay_) {
            return;
        }

        ON_KEY("no_remote_log_file", NoRemoteLogFile_) {
            return;
        }

        ON_KEY("level_compress_file", LevelCompressFile_) {
            return;
        }

        ON_KEY("uaas_mode", UaasMode_) {
            return;
        }

        ON_KEY("queue_limit", QueueLimit_) {
            return;
        }

        if (key == "remote_log_storage") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(RemoteLogStorage_);
            return;
        }

        Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
        return;
    } END_PARSE

private:
    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        return MakeHolder<TTls>(RemoteLogStorage_.Get(), NoRemoteLogFile_, LevelCompressFile_, QueueLimit_, process,
            TSharedCounter(InQueue_, process->WorkerId()), TSharedCounter(Dropped_, process->WorkerId()));
    }

    static const THttpError* ExtractHttpError(const TError& error) noexcept {
        if (const auto* backendError = error.GetAs<TBackendError>()) {
            return backendError->InnerError().GetAs<THttpError>();
        }
        return nullptr;
    }

    class THeaderProcessor: public NExpCommon::IHeaderProcessor {
    public:
        TMaybe<TString> Process(TStringBuf header, TStringBuf value) noexcept override {
            if (Match(HEADER_WHITELIST_REGEXP, header)) {
                return ToString(value);
            }
            if (Match(HEADER_ONLY_LENGTH_REGEXP, header)) {
                return TStringBuilder() << '#' << value.size();
            }
            return Nothing();
        }
    };


    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (tls.RemoteLoggerProps.NoRemoteLogFileExists()) {
            Y_PROPAGATE_ERROR(Submodule_->Run(descr));
        } else {
            // we can use it so we don't create a json string on every request
            // only for the one we interested in
            bool skipJsonRecord = Throttled_ &&  tls.NextRequest > Now();

            TChunksOutputStream out;
            // USEREXP-8329
            const size_t backendErrorsBeforeProcessing = descr.Properties->ConnStats.BackendError;
            NExpCommon::TJsonNoExceptWriter writer(&out);

            if (!skipJsonRecord) {
                THeaderProcessor headerProcessor;
                writer.Open();
                writer.WriteRequestLine(descr.Request->RequestLine());
                writer.WriteAddr(PrintHostAndPort(descr.RemoteAddr()));
                writer.WriteProcessedHeaders(descr.Request->Headers(), &headerProcessor);
                size_t headersTotalLength = 0;

                for(const auto& hdr : descr.Request->Headers()){
                    for(const auto& hdrVal : hdr.second){
                        headersTotalLength += hdr.first.size() + hdrVal.size() + 4; // len(Key + ": " + Value + "\r\n")
                    }
                }
                writer.WriteHeadersLength(headersTotalLength);
                writer.WriteTimestamp(descr.Properties->Start.Seconds());
                if (UaasMode_) {
                    writer.WriteUaasModeFlag();
                }
            }

            NExpCommon::TRemoteLogOutput output{ descr.Output };

            const TInstant start = TInstant::Now();
            Y_TRY(TError, error) {
                return Submodule_->Run(descr.CopyOut(output));
            } Y_CATCH {
                if (!skipJsonRecord) {
                    if (const THttpError* httpError = ExtractHttpError(error)) {
                        writer.WriteStatus(httpError->Code());
                    }
                    const size_t backendErrorsAfterProcessing = descr.Properties->ConnStats.BackendError;
                    writer.WriteBackendErrors(backendErrorsAfterProcessing - backendErrorsBeforeProcessing);
                    writer.WriteResponseTime(TInstant::Now() - start);
                    writer.Close();
                    AddToQueue(tls, out, writer);
                }

                return error;
            }

            if (!skipJsonRecord) {
                const size_t backendErrorsAfterProcessing = descr.Properties->ConnStats.BackendError;
                writer.WriteBackendErrors(backendErrorsAfterProcessing - backendErrorsBeforeProcessing);
                writer.WriteStatus(output.GetStatusCode());
                writer.WriteResponseTime(TInstant::Now() - start);
                writer.Close();
                AddToQueue(tls, out, writer);
            }
        }
        return {};
    }

    void AddToQueue(TTls& tls,TChunksOutputStream& out, NExpCommon::TJsonNoExceptWriter& writer) const {
        if (Throttled_) {
            if (tls.NextRequest > Now()) {
                return;
            }

            tls.NextRequest = Now() + Delay_;
        }

        if (!!tls.RemoteLogger) {
            tls.RemoteLogger->Add(writer.HasExcept(), std::move(out.Chunks()));
        }
    }


private:
    THolder<IModule> RemoteLogStorage_;
    TString NoRemoteLogFile_;
    TString LevelCompressFile_;
    TSharedCounter InQueue_;
    TSharedCounter Dropped_;
    size_t QueueLimit_ = 65536;
    bool UaasMode_ = false;
    TDuration Delay_ = TDuration::Zero();
    bool Throttled_ = false;

};

IModuleHandle* NModRemoteLog::Handle() {
    return TModule::Handle();
}
