#include "module.h"

#include <balancer/modules/log_headers/cookie_masker.h>

#include <balancer/kernel/cookie/cookie.h>
#include <balancer/kernel/cookie/obfuscation.h>
#include <balancer/kernel/helpers/errors.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 <library/cpp/digest/md5/md5.h>

#include <util/generic/string.h>
#include <util/string/split.h>
#include <util/stream/str.h>
#include <util/string/cast.h>


using namespace NConfig;
using namespace NSrvKernel;
using namespace NModLogHeaders;

MODULE_BASE(log_headers, TModuleWithSubModule) {
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

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

        if (LogCookieMeta_) {
            LogSetCookieMeta_ = true;
        }

        if (!HeaderNameFsm_
            && !ResponseHeaderNameFsm_
            && CookieFields_.empty()
            && !LogResponseBody_
            && !LogCookieMeta_
            && !LogSetCookieMeta_
        ) {
            ythrow TConfigParseError() << "neither fsm nor cookie fields nor log_response_body_md5 are configured";
        }
    }

    START_PARSE {
        TString re;

        ON_KEY("name_re", re) {
            HeaderNameFsm_.ConstructInPlace(re, TFsm::TOptions().SetCaseInsensitive(true));
            return;
        }

        TString repsonse_re;
        ON_KEY("response_name_re", repsonse_re) {
            ResponseHeaderNameFsm_.ConstructInPlace(repsonse_re, TFsm::TOptions().SetCaseInsensitive(true));
            return;
        }

        ON_KEY("log_response_body_md5", LogResponseBody_) {
            return;
        }

        ON_KEY("log_cookie_meta", LogCookieMeta_) {
            return;
        }

        ON_KEY("log_set_cookie_meta", LogSetCookieMeta_) {
            return;
        }

        ON_KEY("log_set_cookie", LogSetCookie_) {
            return;
        }

        if (key == "cookie_fields") {
            try {
                const auto cookieFields = value->AsString();
                for (const auto& i: StringSplitter(cookieFields).Split(',')) {
                    CookieFields_.emplace(ToString(i.Token()));
                }
                return;
            } catch (...) {
                ythrow TConfigParseError() << "error parsing " << key.Quote() << ": " << CurrentExceptionMessage();
            }
        }

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

private:
    TError DoRun(const TConnDescr& descr) const noexcept override {
        if (!(LogResponseBody_
            || ResponseHeaderNameFsm_
            || LogSetCookieMeta_
            || LogSetCookie_ && CookieFields_
        )) {
            return PerformRequest(descr);
        }

        MD5 md5;
        TStringStream responseHeadersSB;

        auto output = MakeHttpOutput([&](TResponse&& response, bool forceClose, TInstant deadline) {
            if (ResponseHeaderNameFsm_) {
                if (auto it = response.Headers().FindValues(*ResponseHeaderNameFsm_); it != response.Headers().end()) {
                    for (const auto& value : it->second) {
                        responseHeadersSB << " <::" << it->first << ':' << value << "::>";
                    }
                }
            }

            if (LogSetCookieMeta_ || LogSetCookie_ && CookieFields_) {
                if (const auto it = response.Headers().FindValues("set-cookie"); it != response.Headers().end()) {
                    for (const auto& headerValue : it->second) {
                        TStringBuf nv = StripStringLeft(headerValue.AsStringBuf()).Before(';');
                        TStringBuf n, v;
                        if (nv.TrySplit('=', n, v) && CookieFields_.contains(n)) {
                            // No point in obfuscating the things we already log
                            responseHeadersSB.Str()
                                .append(TStringBuf(" <::SetCookie:"))
                                .append(headerValue.AsStringBuf())
                                .append(TStringBuf("::>"));
                        } else if (LogSetCookieMeta_) {
                            responseHeadersSB.Str().append(TStringBuf(" <::SetCookieMeta:"));
                            ObfuscateSetCookieValue(responseHeadersSB.Str(), headerValue.AsStringBuf());
                            responseHeadersSB.Str().append(TStringBuf("::>"));
                        }
                    }
                }
            }

            return descr.Output->SendHead(std::move(response), forceClose, deadline);
        }, [&](TChunkList lst, TInstant deadline) {
            if (LogResponseBody_) {
                for (auto it = lst.ChunksBegin(); it != lst.ChunksEnd(); ++it) {
                    md5 = md5.Update(it->AsStringBuf());
                }
            }

            return descr.Output->Send(std::move(lst), deadline);
        }, [&](THeaders&& trailers, TInstant deadline) {
            return descr.Output->SendTrailers(std::move(trailers), deadline);
        });

        Y_PROPAGATE_ERROR(PerformRequest(descr.CopyOut(output)));

        if (LogResponseBody_) {
            char rs[33];
            descr.ExtraAccessLog << " <::body_md5:" << md5.End(rs) << "::>";
        }

        if (responseHeadersSB) {
            descr.ExtraAccessLog << responseHeadersSB.Str();
        }

        return {};
    }

    void DoLogCookieHeader(const TAccessLogOutput& logger, const TStringBuf& fields) const noexcept {
        logger << " <::Cookie:" << fields << "::>";
    }

    void LogCookieFieldsOrMeta(const TAccessLogOutput& logger, const TStringStorage& fields) const noexcept {
        Y_ASSERT(!CookieFields_.empty() || LogCookieMeta_); // check before call

        TString cookie;
        TString maskedValue;
        TString cookieMeta;
        if (LogCookieMeta_) {
            cookieMeta.reserve(fields.size());
        }

        EveryCookieKV([&](TNameValue nv) noexcept {
            if (nv.Name && CookieFields_.contains(*nv.Name)) {
                if (!MaskCookieHelper_.ShouldMask(*nv.Name)) {
                    AppendCookieKV(cookie, nv);
                } else {
                    maskedValue = nv.Value;
                    MaskCookieHelper_.MaskCookieParam(maskedValue);
                    AppendCookieC(cookie, *nv.Name + ("=" + maskedValue));
                }
            }
            if (LogCookieMeta_) {
                if (nv.Name) {
                    cookieMeta.append(*nv.Name).append(TStringBuf("=;"));
                } else {
                    cookieMeta.append(nv.Value).append(';');
                }
            }
            }, fields.AsStringBuf()
        );

        if (cookie) {
            DoLogCookieHeader(logger, cookie);
        }

        if (cookieMeta) {
            logger << " <::CookieMeta:"sv << fields.size() << ':' << cookieMeta << "::>"sv;
        }
    }

    void LogCookieHeader(const TAccessLogOutput& logger, const TStringBuf& headerName,
                         const TVector<TStringStorage>& headerValue) const noexcept {
        if (HeaderNameFsm_ && Match(*HeaderNameFsm_, headerName)) {
            for (const auto& fields: headerValue) {
                TString cookie;
                TString maskedValue;
                EveryCookieKV([&](TNameValue nv) noexcept {
                    if (nv.Name) {
                        if (!MaskCookieHelper_.ShouldMask(*nv.Name)) {
                            AppendCookieKV(cookie, nv);
                        } else {
                            maskedValue = nv.Value;
                            MaskCookieHelper_.MaskCookieParam(maskedValue);
                            AppendCookieC(cookie, *nv.Name + ("=" + maskedValue));
                        }
                    }
                    }, fields.AsStringBuf()
                );
                DoLogCookieHeader(logger, cookie);
            }
        }

        // We log these cases even if user has already logged whole cookie header,
        // not sure if it is safe to fix this logic
        if (!CookieFields_.empty() || LogCookieMeta_) {
            for (const auto& fields: headerValue) {
                LogCookieFieldsOrMeta(logger, fields);
            }
        }
    }

    TError PerformRequest(const TConnDescr& descr) const noexcept {
        auto& headers = descr.Request->Headers();

        auto cookieIter = headers.FindValues("cookie"sv);
        if (cookieIter != headers.end()) { // special handling for cookie header
            LogCookieHeader(descr.ExtraAccessLog, cookieIter->first.AsStringBuf(), cookieIter->second);
        }

        if (HeaderNameFsm_) {
            for (auto iter = headers.begin(); iter != headers.end(); ++iter) {
                if (cookieIter != iter && Match(*HeaderNameFsm_, iter->first.AsStringBuf())) {
                    for (const auto& fields: iter->second) {
                        descr.ExtraAccessLog << " <::" << iter->first << ':' << fields << "::>";
                    }
                }
            }
        }

        return Submodule_->Run(descr);
    }

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

    bool DoCanWorkWithoutHTTP() const override {
        return false;
    }

private:
    TMaybe<TFsm> HeaderNameFsm_;
    TMaybe<TFsm> ResponseHeaderNameFsm_;
    THashSet<TString> CookieFields_;
    TCookieMasker MaskCookieHelper_;
    bool LogResponseBody_ = false;
    bool LogCookieMeta_ = false;
    bool LogSetCookieMeta_ = false;
    bool LogSetCookie_ = false;
};

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