#include <balancer/serval/contrib/cone/mun.h>
#include <balancer/serval/core/config.h>

#include <util/stream/file.h>
#include <util/string/cast.h>

static NSv::TAction Const(const YAML::Node& args, NSv::TAuxData&) {
    auto insertHeader = [](NSv::THeaderVector& target, const YAML::Node& k, const YAML::Node& v) {
        const auto& str = k.Scalar();
        CHECK_NODE(k, k.IsScalar(), "header names should be strings");
        CHECK_NODE(k, str.size(), "header names cannot be empty");
        CHECK_NODE(k, str[0] != ':', "cannot set pseudo-headers");
        CHECK_NODE(k, std::find_if(str.begin(), str.end(), isupper) == str.end(), "header names should be lowercase");
        CHECK_NODE(k, str != "content-length", "`content-length` is derived from payload");
        CHECK_NODE(k, str != "transfer-encoding", "this header is forbidden");
        if (!v.IsSequence()) {
            CHECK_NODE(v, v.IsScalar(), "header values must be strings or string sequences");
            target.insert({str, v.Scalar()});
        } else for (const auto& x : v) {
            CHECK_NODE(x, x.IsScalar(), "header values must be strings or string sequences");
            target.insert({str, x.Scalar()});
        }
    };

    NSv::THead head(0, {});
    NSv::THeaderVector tail;
    CHECK_NODE(args, args.IsMap(), "`const` requires an argument");
    for (auto it = args.begin(); ++it != args.end();) {
        CHECK_NODE(it->first, it->first.IsScalar(), "header names must be strings");
        if (it->first.Scalar() == ":code") {
            CHECK_NODE(it->first, head.Code == 0, ":code defined multiple times");
            head.Code = NSv::Required<int>(it->second, [](int c) {
                return c >= 200;
            });
        } else if (it->first.Scalar() == ":tail") {
            CHECK_NODE(it->second, it->second.IsMap(), "should be a map of headers");
            for (const auto& trailer : it->second) {
                insertHeader(tail, trailer.first, trailer.second);
            }
        } else {
            insertHeader(head, it->first, it->second);
        }
    }
    if (head.Code == 0) {
        head.Code = 200;
    }

    const bool payload = (head.Code != 204 && head.Code != 304);
    auto arg = args.begin()->second;
    CHECK_NODE(arg, arg.IsScalar(), "`const` body must be a string");
    CHECK_NODE(arg, payload || arg.Scalar().empty(), "for `const` 204 or 304, the body must be empty");
    auto body = TString(arg.Scalar());
    if (arg.Tag() == "!file") {
        try {
            body = TFileInput(body).ReadAll();
        } catch (TFileError& e) {
            FAIL_NODE(arg, e.what());
        }
    }
    TString contentLength;
    if (payload) {
        head.emplace("content-length", contentLength = ToString(body.size()));
    }

    return [=, args = args, cl = std::move(contentLength)](NSv::IStreamPtr& s) {
        return s->WriteHead(NSv::THead(head)) && (!body || s->Write(body)) && s->Close(NSv::THeaderVector(tail));
    };
}

SV_DEFINE_ACTION("const", Const);
