#include <balancer/serval/contrib/cone/mun.h>
#include <balancer/serval/core/buffer.h>
#include <balancer/serval/core/config.h>
#include <balancer/serval/core/http.h>
#include <balancer/serval/core/io.h>
#include <balancer/serval/core/storage.h>
#include <balancer/serval/mod/restgrpc/restgrpc.ev.pb.h>

#include <google/protobuf/stubs/status.h>
#include <google/protobuf/util/json_util.h>

#include <library/cpp/json/json_writer.h>
#include <library/cpp/protobuf/yql/descriptor.h>
#include <library/cpp/string_utils/base64/base64.h>

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

namespace NSv::NGrpc {
    using TMessageConstructor = NSv::TFunction<THolder<NProtoBuf::Message>()>;
    struct TGrpcAction {
        std::shared_ptr<TMessageConstructor> Input;
        std::shared_ptr<TMessageConstructor> Output;
        TMaybe<TString> DumpDefaultHeader;
        bool dumpOn501 = false;
    };

    using EGrpcCode = google::protobuf::util::StatusCode;
    struct TGrpcStatus {
        EGrpcCode Code;
        TString Message;

        static TGrpcStatus FromProto(const google::protobuf::util::Status& rawStatus) {
            auto msg = google::protobuf::string(rawStatus.message());
            return {rawStatus.code(), ToString(msg)};
        }

        bool Ok() const {
            return Code == EGrpcCode::kOk;
        }
    };

    static bool RespondError(NSv::IStreamPtr req, const TGrpcStatus& status, THeaderVector headers = {}) {
        static TVector<int> httpCodes = {
            200, 499, 500, 400, 504, 404, 409, 403, 429, 400,
            409, 400, 501, 500, 503, 500, 401
        };

        size_t code = static_cast<size_t>(status.Code);
        TString codeString = ToString(code);
        int httpCode = code < httpCodes.size() ? httpCodes[code] : 500;

        headers.erase_if([&] (auto& header) {
            return Interned(header.first, "content-type", "grpc-message", "grpc-status");
        });
        headers.emplace("content-type", "application/json");
        headers.emplace("grpc-message", status.Message);
        headers.emplace("grpc-status", codeString);

        return req->WriteHead({httpCode, std::move(headers)}) && req->Write("{}") && req->Close();
    }

    static TStringBuf EncodeMessage(const NProtoBuf::MessageLite& msg, TString& out) {
        auto delta = msg.ByteSizeLong();
        const ui8 prefix[] = {0, (ui8)(delta >> 24), (ui8)(delta >> 16), (ui8)(delta >> 8), (ui8)delta};
        out.append((const char*)prefix, (const char*)(prefix + 5));
        Y_PROTOBUF_SUPPRESS_NODISCARD msg.AppendPartialToString(&out);
        return out;
    }

    static bool DecodeMessage(TStringBuf data, NProtoBuf::Message* msg) {
        if (data.size() < 5) {
            return false;
        }
        size_t size = (ui32)(ui8)(data)[1] << 24 | (ui32)(ui8)(data)[2] << 16
                    | (ui32)(ui8)(data)[3] << 8  | (ui32)(ui8)(data)[4];

        if (data.size() + 5 < size) {
            return false;
        }

        google::protobuf::io::CodedInputStream input(reinterpret_cast<const ui8*>(data.data() + 5), size);
        Y_PROTOBUF_SUPPRESS_NODISCARD msg->MergePartialFromCodedStream(&input);
        return true;
    }

    static bool IsAllowedHeader(TStringBuf key) {
        return !Interned(key, "content-length", "content-type");
    }

    static bool IsAllowedRequestHeader(TStringBuf key) {
        return IsAllowedHeader(key) && !Interned(key, ":authority", ":scheme", "te");
    }

    static bool IsAllowedResponseHeader(TStringBuf key) {
        return IsAllowedHeader(key);
    }

    static inline TGrpcStatus WrapStream(IStreamPtr& stream, TStringBuf action,
                                            const TGrpcAction& actionInfo, bool dumpDefault)
    {
        auto input = (*actionInfo.Input)();

        TMaybe<TString> inputJson = NSv::ReadFrom(*stream);
        if (!inputJson) {
            return {EGrpcCode::kUnknown, "Empty input!"};
        }
        if (*inputJson) {
            google::protobuf::util::JsonParseOptions options;
            options.ignore_unknown_fields = true;

            auto status = google::protobuf::util::JsonStringToMessage(*inputJson, &*input, options);
            if (!status.ok()) {
                return TGrpcStatus::FromProto(status);
            }
        }
        TString buffer;
        NSv::NGrpc::EncodeMessage(*input, buffer);

        auto headers = std::make_shared<THeaderVector>();
        auto resp = std::make_shared<TString>();
        auto nextStream = NSv::ConstRequestStream(
            {
                "POST",
                action,
                {
                    {":authority", "unknown"},
                    {":scheme", "http"},
                    {"content-type", "application/grpc"},
                    {"te", "trailers"}
                }
            },
            stream->Retain(buffer),
            [headers] (THead& h) {
                headers->insert(h.begin(), h.end());
                return true;
            },
            [resp] (TStringBuf d) {
                resp->append(d);
                return true;
            },
            [headers, resp, stream, outputConstructor = *actionInfo.Output, dumpDefault] (THeaderVector& t) {
                headers->insert(t.begin(), t.end());
                headers->erase_if([] (auto& header) {
                    return !IsAllowedResponseHeader(header.first);
                });
                headers->emplace("content-type", "application/json");
                headers->emplace("access-control-allow-origin", "*");
                headers->emplace("access-control-allow-headers", "*");
                headers->emplace("access-control-allow-methods", "*");

                TGrpcStatus status = {EGrpcCode::kOk, ""};
                for (const auto& header : *headers) {
                    int grpcCode;
                    if (header.first == "grpc-status" && TryFromString(header.second, grpcCode)) {
                        status.Code = static_cast<EGrpcCode>(grpcCode);
                    } else if (header.first == "grpc-message") {
                        status.Message = ToString(header.second);
                    }
                }
                if (!status.Ok()) {
                    return RespondError(stream, status, std::move(*headers));
                }

                TString buffer;
                auto msg = outputConstructor();
                if (DecodeMessage(*resp, &*msg)) {
                    google::protobuf::util::JsonOptions options;
                    options.always_print_primitive_fields = dumpDefault;

                    auto serializeStatus = google::protobuf::util::MessageToJsonString(*msg, &buffer, options);
                    if (!serializeStatus.ok()) {
                        return RespondError(stream, TGrpcStatus::FromProto(serializeStatus), std::move(*headers));
                    }
                } else {
                    return RespondError(stream, {EGrpcCode::kInvalidArgument, "Parse response error!"}, std::move(*headers));
                }
                return stream->WriteHead({200, std::move(*headers)}) && stream->Write(buffer) && stream->Close();
            },
            stream->Log().Fork<NSv::NEv::TGrpcFromRest>(ToString(action)),
            stream->Peer()
            // stream->Stats()
        );
        for (const auto& header : *stream->Head()) {
            if (IsAllowedRequestHeader(header.first)) {
                nextStream->Head()->insert(header);
            }
        }
        stream = nextStream;
        return {EGrpcCode::kOk, ""};
    }

    static std::shared_ptr<TMessageConstructor> MakeMessageConstructor(const YAML::Node& node) {
        TString s = NSv::Required<TString>(node);
        return NSv::StaticData(s, [&] () -> TMessageConstructor {
            return [info = TDynamicInfo::Create(Base64Decode(s))] {
                return THolder<NProtoBuf::Message>(info->Parse("").Release());
            };
        });
    }

    static NSv::TAction RestGrpc(const YAML::Node& node, NSv::TAuxData&) {
        THashMap<TString, TGrpcAction> actions;

        bool dumpOn501 = true;

        for (auto ptr = node.begin(); ++ptr != node.end();) {
            TString name = NSv::Required<TString>(ptr->first);
            CHECK_NODE(ptr->second, ptr->second.IsMap(), "grpc action info must be map");

            TMaybe<TString> dumpDefaultHeader = Nothing();
            if (ptr->second["dump-default-header"].IsDefined()) {
                dumpDefaultHeader = ptr->second["dump-default-header"].as<TString>();
                dumpDefaultHeader->to_lower();
            }

            if (ptr->second["dump-on-501"].IsDefined()) {
                if (!ptr->second["dump-on-501"].as<bool>()) {
                    dumpOn501 = false;
                }
            }

            actions.emplace(name, TGrpcAction{
                MakeMessageConstructor(ptr->second["input"]),
                MakeMessageConstructor(ptr->second["output"]),
                dumpDefaultHeader,
                dumpOn501
            });
        }

        return [actions = std::move(actions), dumpOn501] (NSv::IStreamPtr& req) {
            THead* rqh = req->Head();
            if (!rqh) {
                return false;
            }
            if (size_t pos = rqh->Path().rfind("/")) {
                pos = rqh->Path().rfind("/", pos - 1);
                TStringBuf method = rqh->Path().SubStr(pos);
                if (auto* action = actions.FindPtr(method)) {
                    bool dumpDefault = false;
                    if (action->DumpDefaultHeader) {
                        auto dumpDefaultHeaderPtr = rqh->find(*action->DumpDefaultHeader);
                        if (dumpDefaultHeaderPtr != rqh->end() && dumpDefaultHeaderPtr->second != "0") {
                            dumpDefault = true;
                        }
                    }

                    TGrpcStatus status = WrapStream(req, method, *action, dumpDefault);
                    if ((!status.Ok() MUN_RETHROW) && mun_errno == ECANCELED) {
                        return false;
                    }
                    return status.Ok() ? true : RespondError(req, status);
                }
            }
            if (dumpOn501) {
                Cerr << "=== 501 ===" << Endl;
                Cerr << "Path: " << rqh->Path() << Endl;
                for (auto& action : actions) {
                    Cerr << "Action" << Endl;
                    Cerr << "Key: " << action.first << Endl;
                    Cerr << "Input: " << (*action.second.Input)()->DebugString() << Endl;
                    Cerr << "Output: " << (*action.second.Output)()->DebugString() << Endl;
                }
                Cerr << "===========" << Endl;
            }
            return RespondError(req, {EGrpcCode::kUnimplemented, "Method not found!"});
        };
    }

    SV_DEFINE_ACTION("restgrpc", RestGrpc);
}
