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

#include <util/string/type.h>

static inline TStringBuf IsAntirobotHeader(TStringBuf name) {
    return NSv::Interned(name,
        "x-yandex-internal-request",
        "x-yandex-suspected-robot");
}

static inline TString AsH1(const NSv::THead& h) {
    TStringBuilder ret;
    ret << h.Method << " " << h.PathWithQuery << " HTTP/1.1\r\n";
    for (const auto& header : h) {
        if (header.first == ":authority") {
            ret << "host: " << header.second << "\r\n";
        } else if (!header.first.StartsWith(":")) {
            ret << header.first << ": " << header.second << "\r\n";
        }
    }
    ret << "\r\n";
    return ret;
}

static NSv::TAction Antirobot(const YAML::Node& args, NSv::TAuxData& aux) {
    // Antirobot consumes a part of the payload. After reading it from the request stream,
    // we'll have to put it back.
    class TUnreadStream : public NSv::TStreamProxy {
    public:
        TUnreadStream(NSv::IStreamPtr parent, TString data)
            : NSv::TStreamProxy(std::move(parent))
            , Buffer_(data)
        {
        }

        TMaybe<TStringBuf> Read() noexcept override {
            if (Pos_ < Buffer_.size()) {
                return TStringBuf(Buffer_).Skip(Pos_);
            }
            return TStreamProxy::Read();
        }

        bool Consume(size_t n) noexcept override {
            if (Pos_ + n < Buffer_.size()) {
                Pos_ += n;
                return true;
            }
            n -= Buffer_.size() - Pos_;
            Pos_ = Buffer_.size();
            return TStreamProxy::Consume(n);
        }

        bool AtEnd() const noexcept override {
            return Pos_ == Buffer_.size() && NSv::TStreamProxy::AtEnd();
        }

        size_t Replay() noexcept override {
            auto pos = NSv::TStreamProxy::Replay();
            if (pos > Buffer_.size()) {
                return pos;
            }
            Pos_ = 0;
            return 0;
        }

    private:
        TString Buffer_;
        size_t Pos_ = 0;
    };

    CHECK_NODE(args, args.IsMap(), "`antirobot` requires an argument");
    auto proxy = aux.Action(args.begin()->second);
    auto limit = NSv::Optional<size_t>(args["cut-payload"], 512);
    return [=, &captchas = aux.Signal("captchas_dmmm")](NSv::IStreamPtr& req) {
        auto rqh = req->Head();
        if (!rqh MUN_RETHROW) {
            return false;
        }
        auto payload = AsH1(*rqh);
        if (limit) {
            // FIXME this requires h100 at some point before; should handle 100-continue automatically
            auto chunk = NSv::ReadFrom(*req, limit);
            if (!chunk MUN_RETHROW) {
                return false;
            }
            payload += *chunk;
            if (*chunk) {
                req = std::make_shared<TUnreadStream>(std::move(req), std::move(*chunk));
            }
        }
        rqh->erase_if([](auto& h) {
            return !!IsAntirobotHeader(h.first);
        });
        bool forward = false;
        NSv::IStreamPtr fake = NSv::ConstRequestStream(
            {"POST", "/fullreq", {{":authority", "unknown"}, {":scheme", "http"}}}, payload,
            [&](NSv::THead& rsp) {
                auto it = rsp.find("x-forwardtouser-y");
                if (it != rsp.end() && ::IsTrue(it->second)) {
                    forward = true;
                    captchas++;
                    rsp.erase_if([](auto& h) {
                        return IsAntirobotHeader(h.first) || h.first == "x-forwardtouser-y";
                    });
                    return req->WriteHead(rsp);
                }
                for (const auto& h : rsp) {
                    if (auto interned = IsAntirobotHeader(h.first)) {
                        rqh->emplace(interned, ::IsTrue(h.second) ? TStringBuf("1") : TStringBuf("0"));
                    }
                }
                return true;
            },
            [&](TStringBuf data) {
                return !forward || req->Write(data);
            },
            [&](NSv::THeaderVector& tail) {
                return !forward || req->Close(tail);
            }
        );
        // If not forwarding, fail only on client disconnect.
        return proxy(fake) || (!forward && mun_errno != ECANCELED);
    };
}

SV_DEFINE_ACTION("antirobot", Antirobot);
