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

#include <util/random/random.h>

static NSv::TAction ForceClose(const YAML::Node& args, NSv::TAuxData&) {
    class TAlwaysClosingStream: public NSv::TStreamProxy {
    public:
        using NSv::TStreamProxy::TStreamProxy;

        bool WriteHead(NSv::THead& head) noexcept override {
            auto connHeader = head.find("connection");
            if (!head.IsInformational() && (connHeader == head.end() || connHeader->second != "close")) {
                head.emplace("connection", "close");
            }
            return NSv::TStreamProxy::WriteHead(head);
        }
    };

    double threshold = 1.0;
    if (args.IsMap()) {
        threshold = NSv::Required<double>(args.begin()->second, [](double p) {
            return 0 < p && p <= 1;
        });
    }
    return [=](NSv::IStreamPtr& s) {
        if (RandomNumber<double>() < threshold) {
            s = std::make_shared<TAlwaysClosingStream>(std::move(s));
        }
        return true;
    };
}

SV_DEFINE_ACTION("force-close", ForceClose);
