#include <balancer/serval/core/config.h>

static inline TStringBuf IsExpHeader(TStringBuf name) {
    return NSv::Interned(name,
        "x-yandex-logstatuid",
        "x-yandex-expboxes",
        "x-yandex-expboxes-pre",
        "x-yandex-expconfigversion",
        "x-yandex-expconfigversion-pre",
        "x-yandex-expflags",
        "x-yandex-expflags-pre",
        "x-yandex-expsplitparams",
        "x-yandex-is-staff-login");
}

static inline bool HasAnyExpHeader(const NSv::THead& h) {
    for (const auto& hdr : h) {
        if (IsExpHeader(hdr.first)) {
            return true;
        }
    }
    return false;
}

static NSv::TAction Experiments(const YAML::Node& args, NSv::TAuxData& aux) {
    CHECK_NODE(args, args.IsMap(), "`experiments` requires an argument");
    auto proxy = aux.Action(args.begin()->second);
    auto trusted = NSv::Optional(args["trusted"], false);
    return [=, &forwarded = aux.Signal("trusted_dmmm")](NSv::IStreamPtr& req) {
        auto rqh = req->Head();
        if (!rqh MUN_RETHROW) {
            return false;
        }
        if (trusted && HasAnyExpHeader(*rqh)) {
            forwarded++;
            return true;
        }
        rqh->erase_if([](auto& h) {
            return !!IsExpHeader(h.first);
        });
        NSv::THead head = *rqh;
        head.erase_if([](auto& h) {
            return EqualToOneOf(h.first, "content-length", "transfer-encoding", "connection");
        });
        NSv::IStreamPtr fake = NSv::ConstRequestStream(std::move(head), {},
            [&](NSv::THead& rsp) {
                for (const auto& h : rsp) {
                    if (auto interned = IsExpHeader(h.first)) {
                        rqh->emplace(interned, req->Retain(TString(h.second)));
                    }
                }
                return true;
            },
            [&](TStringBuf) {
                return true;
            },
            [&](NSv::THeaderVector&) {
                return true;
            }
        );
        return proxy(fake) || mun_errno != ECANCELED;
    };
}

SV_DEFINE_ACTION("experiments", Experiments);
