#include "module.h"

#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_pire.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/log/errorlog.h>

#include <util/datetime/base.h>
#include <util/generic/singleton.h>
#include <util/random/random.h>
#include <util/string/builder.h>
#include <util/system/compiler.h>


using namespace NConfig;
using namespace NSrvKernel;

namespace {
    class TSaveStatusOutput final : public IHttpOutput {
    public:
        explicit TSaveStatusOutput(IHttpOutput& output)
            : Output_(output)
        {}
        ui16 StatusCode() const noexcept {
            return StatusCode_;
        }
    private:
        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            if (TError error = Output_.Send(std::move(lst), deadline)) {
                StatusCode_ = 0;
                return error;
            }
            return {};
        }
        TError DoSendHead(TResponse&& response, const bool forceClose, TInstant deadline) override {
            StatusCode_ = response.ResponseLine().StatusCode;
            if (TError error = Output_.SendHead(std::move(response), forceClose, deadline)) {
                StatusCode_ = 0;
                return error;
            }
            return {};
        }
        TError DoSendTrailers(THeaders&& trailers, TInstant deadline) override {
            if (TError error = Output_.SendTrailers(std::move(trailers), deadline)) {
                StatusCode_ = 0;
                return error;
            }
            return {};
        }
        IHttpOutput& Output_;
        ui16 StatusCode_ = 0;
    };
}

Y_TLS(accesslog) {
    TLog* Log = nullptr;
};

MODULE_WITH_TLS_BASE(accesslog, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);
        if (LogName_.empty()) {
            Control->CreateDefaultLog();
        }

        if (AdditionalPortHeaderFsm_ && !AdditionalIpHeaderFsm_) {
            ythrow TConfigParseError() << "\"additional_port_header\" specified, but it requires \"additional_ip_header\"";
        }

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

        if (AdditionalPortHeaderFsm_) {
            UberFsm_.Reset(new TFsm(
                        THostFsm::Instance() // 0
                      | TRefererFsm::Instance() // 1
                      | *AdditionalIpHeaderFsm_ // 2
                      | *AdditionalPortHeaderFsm_ // 3
            ));
        } else if (AdditionalIpHeaderFsm_) {
            UberFsm_.Reset(new TFsm(
                        THostFsm::Instance() // 0
                      | TRefererFsm::Instance() // 1
                      | *AdditionalIpHeaderFsm_ // 2
            ));
        } else {
            UberFsm_.Reset(new TFsm(
                        THostFsm::Instance() // 0
                      | TRefererFsm::Instance() // 1
            ));
        }
    }

private:
    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TTls>();
        tls->Log = process->GetLog(LogName_);
        return tls;
    }

    START_PARSE {
        {
            ON_KEY("log", LogName_) {
                Control->CreateLog(LogName_);
                return;
            }
        }

        if (key == "extended") {
            return;
        }

        {
            bool logInstanceDir = false;
            ON_KEY("log_instance_dir", logInstanceDir) {
                if (logInstanceDir) {
                    InstanceDir_ = TFsPath::Cwd().GetName();
                }
                return;
            }
        }

        {
            TString additionalIpHeader;

            ON_KEY("additional_ip_header", additionalIpHeader) {
                AdditionalIpHeaderFsm_.Reset(new TFsm(additionalIpHeader, TFsm::TOptions().SetCaseInsensitive(true)));
                return;
            }
        }

        {
            TString additionalPortHeader;

            ON_KEY("additional_port_header", additionalPortHeader) {
                AdditionalPortHeaderFsm_.Reset(new TFsm(additionalPortHeader, TFsm::TOptions().SetCaseInsensitive(true)));
                return;
            }
        }

        ON_KEY("errors_only", LogErrorsOnly_) {
            return;
        }

        ON_KEY("logged_share", LoggedShare_) {
            return;
        }

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

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        TTempBufOutput extraAccessLog;

        IOutputStream* accessLogSlave = &extraAccessLog;
        if (!descr.ExtraAccessLog.HasNullSlave()) {
            accessLogSlave = descr.ExtraAccessLog.Slave();
        }

        TAccessLogSummary summary;
        TAccessLogSummary* summaryPtr = descr.ExtraAccessLog.Summary() ?: &summary;

        TSaveStatusOutput out(*descr.Output);

        auto connDescr = descr.Copy();
        connDescr.ExtraAccessLog = { accessLogSlave, summaryPtr };
        connDescr.Output = &out;

        TTempBufOutput* const finalExtraAccessLog = dynamic_cast<TTempBufOutput*>(accessLogSlave);
        Y_VERIFY(finalExtraAccessLog);

        TString hostHeaderStr;
        TString refererHeaderStr;
        TString additionalIpHeaderStr;
        TString additionalPortHeaderStr;

        if (Y_LIKELY(descr.Request)) {
            for (auto& hdr: descr.Request->Headers()) {
                if (hdr.second.empty()) {
                    continue;
                }
                TMatcher matcher{ *UberFsm_ };

                if (Match(matcher, hdr.first.AsStringBuf()).Final()) {
                    switch (*matcher.MatchedRegexps().first) {
                        case 0: // host
                            hostHeaderStr = hdr.second.back().AsStringBuf();
                            break;
                        case 1: // referer
                            refererHeaderStr = hdr.second.back().AsStringBuf();
                            break;
                        case 2: // ip header, probably X-Forwarded-For-Y
                            additionalIpHeaderStr = hdr.second.back().AsStringBuf();
                            break;
                        case 3: // port header, probably X-Source-Port-Y
                            additionalPortHeaderStr = hdr.second.back().AsStringBuf();
                            break;
                        default:
                            break;
                    }
                }
            }
        }

        TString startTimeStr;
        try {
            // storing it here to increase probability of getting it from formatter's cache
            startTimeStr = descr.Properties->Start.ToStringLocal();
        } Y_TRY_STORE(TSystemError, std::length_error, yexception);

        TError error = Submodule_->Run(connDescr);

        if (error || LogCurrentEntry(out.StatusCode())) {
            WriteLog(descr, TStringBuf(finalExtraAccessLog->Data(), finalExtraAccessLog->Filled()),
                     startTimeStr, hostHeaderStr, refererHeaderStr, additionalIpHeaderStr,
                     additionalPortHeaderStr, tls, out.StatusCode(), *summaryPtr);
        }
        return error;
    }

    bool DoCanWorkWithoutHTTP() const noexcept override {
        return true;
    }

private:
    bool LogCurrentEntry(ui16 const statusCode) const noexcept {
        if (LogErrorsOnly_) {
            return statusCode < 200 || statusCode >= 400;
        }
        return LoggedShare_ >= 1.0 || RandomNumber<double>() < LoggedShare_;
    }

    void WriteLog(
        const TConnDescr& descr,
        TStringBuf extraAccessLog,
        const TString& startTimeStr,
        const TString& hostHeader,
        const TString& refererHeader,
        const TString& additionalIpHeader,
        const TString& additionalPortHeader,
        const TTls& tls,
        const ui16 statusCode,
        const TAccessLogSummary& summary)
    const noexcept {
        const auto now = Now();

        TStringBuilder ss;
        ss.reserve(4096);

        ss << descr.RemoteAddrPortStr() << "\t";

        if (AdditionalIpHeaderFsm_) {
            ss << '"' << additionalIpHeader;
            if (AdditionalPortHeaderFsm_) {
                ss << ':' << additionalPortHeader;
            }
            ss << "\"\t";
        }

        ss << startTimeStr << "\t";

        ss << "\"";
        if (descr.Request) {
            const TRequestLine& requestLine = descr.Request->RequestLine();
            requestLine.PrintNoCRLFWithProto(ss, descr.Ssl().NextProto);
        }
        ss << "\"\t";

        ss << now - descr.Properties->Start << "\t";
        ss << "\"" << refererHeader << "\"\t";
        ss  << "\"" << hostHeader << "\"\t";

        if (InstanceDir_) {
            ss << InstanceDir_ << " ";
        }

        ss << " [" << GetHandle()->Name() << " <::status:" << statusCode << "::>";
        Y_ASSERT(summary.AnsweredModule);
        if (summary.AnsweredModule) {
            ss << " <::module:" << summary.AnsweredModule << "::>";
        }
        Y_ASSERT(summary.AnswerReason);
        if (summary.AnswerReason) {
            ss << " <::reason:" << summary.AnswerReason << "::>";
        }
        ss << extraAccessLog;
        ss << "]";

        *tls.Log << ss << Endl;
    }

private:
    TString LogName_;
    THolder<TFsm> AdditionalIpHeaderFsm_;
    THolder<TFsm> AdditionalPortHeaderFsm_;
    THolder<TFsm> UberFsm_;
    TMaybe<TString> InstanceDir_;
    bool LogErrorsOnly_ = false;
    double LoggedShare_ = 1.0;
};

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