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

static inline bool IsLaasHeader(TStringBuf name) {
    return name == "x-laas-answered" || name == "x-ip-properties" || name.StartsWith("x-region-");
}

static inline bool HasRegion(const NSv::THead& h) {
    for (const auto& hdr : h) {
        if (hdr.first.StartsWith("x-region-")) {
            return true;
        }
    }
    return false;
}

static NSv::TAction Geobase(const YAML::Node& args, NSv::TAuxData& aux) {
    CHECK_NODE(args, args.IsMap(), "`geobase` requires an argument");
    auto proxy = aux.Action(args.begin()->second);
    auto trusted = NSv::Optional<bool>(args["trusted"], false);
    auto path = NSv::Optional<TString>(args["path"], "/region?response_format=header&version=1&service=balancer", [](const TString& s) {
        return TStringBuf(s).StartsWith("/");
    });
    return [=,
        &noAddress = aux.Signal("no-ip-header_dmmm"),
        &noRegion  = aux.Signal("no-region_dmmm"),
        &forwarded = aux.Signal("trusted_dmmm")
    ](NSv::IStreamPtr& req) {
        auto rqh = req->Head();
        if (!rqh MUN_RETHROW) {
            return false;
        }
        if (trusted && rqh->find("x-laas-answered") != rqh->end() && HasRegion(*rqh)) {
            forwarded++;
            return true;
        }

        Y_DEFER {
            bool hasRegion = HasRegion(*rqh);
            if (!hasRegion) {
                noRegion++;
            }
            rqh->emplace("x-laas-answered", hasRegion ? "1" : "0");
        };
        rqh->erase_if([](auto& h) { return IsLaasHeader(h.first); });
        if (rqh->find("x-forwarded-for-y") == rqh->end()) {
            noAddress++;
            return true;
        }

        NSv::THead head{"GET", path, *rqh};
        head.erase_if([](auto& h) {
            return EqualToOneOf(h.first, "content-length", "transfer-encoding", "connection", "x-url-prefix");
        });
        auto scheme = rqh->find(":scheme")->second;
        auto host = rqh->find(":authority")->second;
        auto urlPrefix = TString::Join(scheme != "unknown" ? scheme : "http", "://", host, rqh->Path());
        head.emplace("x-url-prefix", urlPrefix);
        NSv::IStreamPtr fake = NSv::ConstRequestStream(std::move(head), {},
            [&](NSv::THead& rsp) {
                for (const auto& h : rsp) {
                    if (IsLaasHeader(h.first)) {
                        rqh->emplace(req->Retain(TString(h.first)), req->Retain(TString(h.second)));
                    }
                }
                return true;
            },
            [&](TStringBuf) {
                return true;
            },
            [&](NSv::THeaderVector&) {
                return true;
            }
        );
        return proxy(fake) || mun_errno != ECANCELED;
    };
}

SV_DEFINE_ACTION("geobase", Geobase);
