XFE.SpaceEngineers.AgentBridge

using System;
using System.Globalization;
using System.IO;
using System.Text;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

namespace XFE.SeAgent.Plugin.Protocol
{
    public sealed class AgentRequest
    {
        public JToken Id { get; }
        public string Method { get; }
        public JObject Parameters { get; }
        public int TimeoutMs { get; }

        internal AgentRequest(JToken id, string method, JObject parameters, int timeoutMs)
        {
            Id = id;
            Method = method;
            Parameters = parameters;
            TimeoutMs = timeoutMs;
        }
    }

    /// <summary>A dispatcher can use this exception for a bounded, public RPC error.</summary>
    public class AgentRpcException : Exception
    {
        public string Code { get; }

        public AgentRpcException(string code, string message) : base(message)
        {
            Code = code;
        }
    }

    internal sealed class AgentProtocolException : AgentRpcException
    {
        internal JToken RequestId { get; }

        internal AgentProtocolException(string code, string message, JToken requestId = null) : base(code, message)
        {
            RequestId = requestId;
        }
    }

    public static class AgentProtocol
    {
        public const int MaxRequestBytes = 1024 * 1024;
        public const int MaxResponseBytes = 4 * 1024 * 1024;
        public const int DefaultTimeoutMs = 10000;
        public const int MaximumTimeoutMs = 30000;
        internal static readonly Encoding Utf8 = new UTF8Encoding(false, true);

        public static AgentRequest ParseRequest(string text)
        {
            if (text == null) throw new ArgumentNullException(nameof(text));
            if (Utf8.GetByteCount(text) > MaxRequestBytes)
                throw new AgentProtocolException("request_too_large", "Request exceeds the 1 MiB limit.");
            JObject body;
            try
            {
                using (var reader = new JsonTextReader(new StringReader(text)))
                {
                    reader.DateParseHandling = DateParseHandling.None;
                    reader.FloatParseHandling = FloatParseHandling.Decimal;
                    reader.Culture = CultureInfo.InvariantCulture;
                    reader.MaxDepth = 32;
                    body = JObject.Load(reader, new JsonLoadSettings { DuplicatePropertyNameHandling = DuplicatePropertyNameHandling.Error });
                    if (reader.Read()) throw new JsonReaderException("Only one JSON object is allowed.");
                }
            }
            catch (JsonException)
            {
                throw new AgentProtocolException("parse_error", "Request must be one JSON object with a maximum depth of 32.");
            }

            JToken id = body["id"];
            if (id == null || (id.Type != JTokenType.String && id.Type != JTokenType.Integer)
                || (id.Type == JTokenType.String ? ((string)id).Length : id.ToString(Formatting.None).Length) > 128)
                throw new AgentProtocolException("invalid_request", "id must be a string or integer of at most 128 characters.");
            id = id.DeepClone();
            JToken methodValue = body["method"];
            string method = methodValue != null && methodValue.Type == JTokenType.String ? (string)methodValue : null;
            if (string.IsNullOrEmpty(method) || method.Length > 128 || !IsMethodName(method))
                throw new AgentProtocolException("invalid_request", "method must contain 1–128 ASCII letters, digits, dots, dashes or underscores.", id);

            JToken arguments = body["params"];
            if (arguments != null && arguments.Type != JTokenType.Null && !(arguments is JObject))
                throw new AgentProtocolException("invalid_params", "params must be an object.", id);

            int timeoutMs = DefaultTimeoutMs;
            JToken timeout = body["timeoutMs"];
            if (timeout != null)
            {
                if (timeout.Type != JTokenType.Integer
                    || !int.TryParse(timeout.ToString(Formatting.None), NumberStyles.None, CultureInfo.InvariantCulture, out timeoutMs)
                    || timeoutMs < 1 || timeoutMs > MaximumTimeoutMs)
                    throw new AgentProtocolException("invalid_request", "timeoutMs must be an integer from 1 to 30000.", id);
            }
            return new AgentRequest(id, method, arguments as JObject ?? new JObject(), timeoutMs);
        }

        private static bool IsMethodName(string method)
        {
            foreach (char value in method)
                if (!(value >= 'a' && value <= 'z') && !(value >= 'A' && value <= 'Z')
                    && !(value >= '0' && value <= '9') && value != '.' && value != '-' && value != '_') return false;
            return true;
        }

        public static JObject Success(JToken id, JObject result)
        {
            return new JObject
            {
                ["id"] = id?.DeepClone() ?? JValue.CreateNull(),
                ["result"] = result?.DeepClone() ?? (JToken)JValue.CreateNull()
            };
        }

        public static JObject Error(JToken id, string code, string message)
        {
            return new JObject
            {
                ["id"] = id?.DeepClone() ?? JValue.CreateNull(),
                ["error"] = new JObject { ["code"] = Limit(code, 64), ["message"] = Limit(message, 2048) }
            };
        }

        private static string Limit(string value, int length)
        {
            return string.IsNullOrEmpty(value) ? "Unknown error." : value.Length <= length ? value : value.Substring(0, length);
        }

        /// <summary>Returns UTF-8 JSON plus LF. Serialization itself is bounded, not only the final allocation.</summary>
        public static byte[] SerializeResponse(JObject response)
        {
            if (response == null) throw new ArgumentNullException(nameof(response));
            using (var output = new LimitedMemoryStream(MaxResponseBytes))
            {
                using (var writer = new StreamWriter(output, Utf8, 4096, true))
                using (var json = new JsonTextWriter(writer) { Formatting = Formatting.None, CloseOutput = false })
                {
                    response.WriteTo(json);
                    json.Flush();
                    writer.Flush();
                }
                output.WriteByte((byte)'\n');
                return output.ToArray();
            }
        }

        private sealed class LimitedMemoryStream : MemoryStream
        {
            private readonly int limit;

            internal LimitedMemoryStream(int limit) : base(4096) { this.limit = limit; }

            public override void Write(byte[] buffer, int offset, int count)
            {
                if (Position + count > limit) throw new AgentRpcException("response_too_large", "Response exceeds the 4 MiB limit. Request fewer items.");
                base.Write(buffer, offset, count);
            }

            public override void WriteByte(byte value)
            {
                if (Position >= limit) throw new AgentRpcException("response_too_large", "Response exceeds the 4 MiB limit. Request fewer items.");
                base.WriteByte(value);
            }
        }
    }
}