mcdiscordbot/MinecraftDiscordBot/ConnectedComputer.cs

using Discord;
using Discord.WebSocket;
using Fleck;
using Newtonsoft.Json;
using System.Diagnostics;
using System.Runtime.Serialization;
using System.Text;

namespace MinecraftDiscordBot;

public delegate Task<TResponse> HandleCommandDelegate<TResponse>(SocketUserMessage message, string[] parameters, CancellationToken ct);
public delegate Task HandleCommandDelegate(SocketUserMessage message, string[] parameters, CancellationToken ct);

[AttributeUsage(AttributeTargets.Method, Inherited = true, AllowMultiple = false)]
public sealed class CommandHandlerAttribute : Attribute {
    public CommandHandlerAttribute(string commandName) => CommandName = commandName;
    public string CommandName { get; }
}

public class ConnectedComputer : CommandRouter<ResponseType> {
    protected readonly IWebSocketConnection _socket;
    public ConnectedComputer(IWebSocketConnection socket) : base() {
        socket.OnMessage = OnMessage;
        _socket = socket;
    }

    private void OnMessage(string message) {
        if (JsonConvert.DeserializeObject<ReplyMessage>(message) is not ReplyMessage msg) return;
        IChunkWaiter? waiter;
        lock (_syncRoot) {
            if (!_waits.TryGetValue(msg.AnswerId, out waiter)) {
                Program.LogWarningAsync("Socket", $"Invalid wait id '{msg.AnswerId}'!");
                return;
            }
        }
        if (!msg.Success) waiter.SetUnsuccessful();
        waiter.AddChunk(msg.Chunk, msg.Total, msg.Result);
        if (waiter.Finished || waiter.IsCancellationRequested)
            lock (_syncRoot)
                _waits.Remove(waiter.ID);
    }
    public Task Send(string message) => _socket.Send(message);
    protected Task Send(Message message) => Send(JsonConvert.SerializeObject(message));
    private readonly object _syncRoot = new();
    private readonly Dictionary<int, IChunkWaiter> _waits = new();
    private readonly Random _rnd = new();
    public IWebSocketConnectionInfo ConnectionInfo => _socket.ConnectionInfo;

    protected interface IChunkWaiter {
        bool Finished { get; }
        int ID { get; }
        bool IsCancellationRequested { get; }
        void AddChunk(int chunkId, int totalChunks, string value);
        void SetUnsuccessful();
    }

    protected class ChunkWaiter<T> : IChunkWaiter {
        public int ID { get; }
        private readonly CancellationToken _ct;
        public ChunkWaiter(int id, Func<string, T> resultParser, CancellationToken ct) {
            ID = id;
            this.resultParser = resultParser;
            _ct = ct;
        }
        private readonly TaskCompletionSource<T> tcs = new();
        private readonly Func<string, T> resultParser;
        public Task<T> Task => tcs.Task.WaitAsync(_ct);
        public bool Finished { get; private set; } = false;
        public bool IsCancellationRequested => _ct.IsCancellationRequested;
        private string?[]? _chunks = null;
        private int _receivedChunks = 0;
        private bool _success = true;
        private readonly object _syncRoot = new();
        public void AddChunk(int chunkId, int totalChunks, string value) {
            lock (_syncRoot) {
                if (_chunks is null) _chunks = new string[totalChunks];
                else if (_chunks.Length != totalChunks) {
                    Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException("Different numbers of chunks in same message ID!"));
                    return;
                }
                ref string? chunk = ref _chunks[chunkId - 1]; // Lua 1-indexed
                if (chunk is not null) {
                    Program.LogErrorAsync(Program.WebSocketSource, new InvalidOperationException($"Chunk with ID {chunkId} was already received!"));
                    return;
                }
                chunk = value;
            }
            if (++_receivedChunks == totalChunks) FinalizeResult(_chunks);
        }
        private void FinalizeResult(string?[] _chunks) {
            var resultString = string.Concat(_chunks);
            if (_success) tcs.SetResult(resultParser(resultString));
            else tcs.SetException(new ReplyException(resultString));
            Finished = true;
        }
        public void SetUnsuccessful() => _success = false;
    }

    protected int GetFreeId() {
        var attempts = 0;
        while (true) {
            var id = _rnd.Next();
            if (!_waits.ContainsKey(id))
                return id;
            Program.LogWarningAsync(Program.WebSocketSource, $"Could not get a free ID after {++attempts} attempts!");
        }
    }

    protected ChunkWaiter<T> GetWaiter<T>(Func<string, T> resultParser, CancellationToken ct) {
        ChunkWaiter<T> waiter;
        lock (_syncRoot) {
            waiter = new ChunkWaiter<T>(GetFreeId(), resultParser, ct);
            _waits.Add(waiter.ID, waiter);
        }
        return waiter;
    }

    protected static Func<string, T> Deserialize<T>() => msg
         => JsonConvert.DeserializeObject<T>(msg) ?? throw new InvalidProgramException("Empty response!");

    public const string Role = "rs";
    private const string CmdEnergyUsage = "energyusage";
    private const string CmdEnergyStorage = "energystorage";
    private const string CmdListItems = "listitems";
    private const string CmdItemName = "itemname";
    private const string CmdListFluids = "listfluids";

    public async Task<int> GetEnergyUsageAsync(CancellationToken ct) {
        var waiter = GetWaiter(int.Parse, ct);
        await Send(new RequestMessage(waiter.ID, CmdEnergyUsage));
        return await waiter.Task;
    }
    public async Task<int> GetEnergyStorageAsync(CancellationToken ct) {
        var waiter = GetWaiter(int.Parse, ct);
        await Send(new RequestMessage(waiter.ID, CmdEnergyStorage));
        return await waiter.Task;
    }
    public async Task<IEnumerable<Item>> ListItemsAsync(CancellationToken ct) {
        var waiter = GetWaiter(Deserialize<IEnumerable<Item>>(), ct);
        await Send(new RequestMessage(waiter.ID, CmdListItems));
        return await waiter.Task;
    }

    public async Task<IEnumerable<Fluid>> ListFluidsAsync(CancellationToken ct) {
        var waiter = GetWaiter(Deserialize<IEnumerable<Fluid>>(), ct);
        await Send(new RequestMessage(waiter.ID, CmdListFluids));
        return await waiter.Task;
    }

    [CommandHandler(CmdEnergyStorage)]
    public async Task<ResponseType> HandleEnergyStorage(SocketUserMessage message, string[] parameters, CancellationToken ct)
        => ResponseType.AsString($"Refined Storage system stores {await GetEnergyStorageAsync(ct)} RF/t");
    [CommandHandler(CmdEnergyUsage)]
    public async Task<ResponseType> HandleEnergyUsage(SocketUserMessage message, string[] parameters, CancellationToken ct)
        => ResponseType.AsString($"Refined Storage system currently uses {await GetEnergyUsageAsync(ct)} RF/t");
    [CommandHandler(CmdItemName)]
    public async Task<ResponseType> HandleItemName(SocketUserMessage message, string[] parameters, CancellationToken ct) {
        if (parameters.Length < 2) return ResponseType.AsString($"Usage: {CmdItemName} filters...");
        else {
            var items = await FilterItems(message, parameters[1..], ct);
            var sb = new StringBuilder();
            sb.AppendLine("Did you mean:");
            sb.AppendJoin("\n", items.Select(i => i.ToString()));
            return ResponseType.AsString(sb.ToString());
        }
    }

    private Task<IEnumerable<Item>> FilterItems(SocketUserMessage message, IEnumerable<string> filters, CancellationToken ct)
        => FilterItems(message, filters.Select(ItemFilter.Parse), ct);

    private async Task<IEnumerable<Item>> FilterItems(SocketUserMessage message, IEnumerable<ItemFilter> filters, CancellationToken ct) {
        var items = Items?.ToList().AsEnumerable();
        if (items is null) items = (await RefreshItemList(ct)).ToList();
        foreach (var filter in filters)
            items = items.Where(filter.MatchItem);
        return items.ToList();
    }

    [CommandHandler(CmdListFluids)]
    public async Task<ResponseType> HandleFluidListing(SocketUserMessage message, string[] parameters, CancellationToken ct) {
        var sb = new StringBuilder();
        sb.Append("The Refined Storage system stores those fluids:");
        var fluids = await ListFluidsAsync(ct);
        foreach (var fluid in fluids.OrderByDescending(i => i.Amount))
            if (fluid.Amount > 10000) sb.AppendFormat("\n{0:n2} B of {1}", fluid.Amount / 1000.0f, fluid.DisplayName);
            else sb.AppendFormat("\n{0:n0} mB of {1}", fluid.Amount, fluid.DisplayName);
        return ResponseType.AsString(sb.ToString());
    }

    private List<Item>? Items;
    private readonly object _itemLock = new();

    [CommandHandler(CmdListItems)]
    public async Task<ResponseType> HandleItemListing(SocketUserMessage message, string[] parameters, CancellationToken ct) {
        var sb = new StringBuilder();
        sb.Append("The Refined Storage system currently stores these items:");
        var items = await RefreshItemList(ct);
        lock (_itemLock) {
            int taken = 0;
            foreach (var item in items) {
                if (sb.Length > 500) break;
                sb.AppendFormat("\n{0:n0}x {1}", item.Amount, item.DisplayName);
                taken++;
            }
            if (items.Count > taken) sb.AppendFormat("\nand {0} more items.", items.Skip(taken).Sum(i => i.Amount));
        }
        return ResponseType.AsString(sb.ToString());
    }

    private async Task<List<Item>> RefreshItemList(CancellationToken ct) {
        var response = await ListItemsAsync(ct);
        lock (_itemLock) {
            Items = response.OrderByDescending(i => i.Amount).ToList();
            return Items;
        }
    }

    public override Task<ResponseType> RootAnswer(SocketUserMessage message, CancellationToken ct)
        => Task.FromResult(ResponseType.AsString("The Minecraft server is connected!"));
    public override Task<ResponseType> FallbackHandler(SocketUserMessage message, string method, string[] parameters, CancellationToken ct)
        => Task.FromResult(ResponseType.AsString($"What the fuck do you mean by '{method}'?"));
}

[Serializable]
public class ReplyException : Exception {
    public ReplyException() { }
    public ReplyException(string message) : base(message) { }
    public ReplyException(string message, Exception inner) : base(message, inner) { }
    protected ReplyException(SerializationInfo info, StreamingContext context) : base(info, context) { }
}

[JsonObject(MemberSerialization.OptIn, Description = "Describes an item in a Refined Storage system.", MissingMemberHandling = MissingMemberHandling.Ignore)]
[DebuggerDisplay($"{{{nameof(ToString)}(),nq}}")]
public class Item : Fluid {
    [JsonProperty("fingerprint", Required = Required.Always)]
    public Md5Hash Fingerprint { get; set; } = default!;
    [JsonProperty("nbt", Required = Required.DisallowNull)]
    public dynamic? NBT { get; set; }
    public override string ToString() => $"{Amount:n0}x {DisplayName}";
}

[JsonObject(MemberSerialization.OptIn, Description = "Describes a fluid in a Refined Storage system.", MissingMemberHandling = MissingMemberHandling.Ignore)]
[DebuggerDisplay($"{{{nameof(ToString)}(),nq}}")]
public class Fluid {
    [JsonProperty("amount", Required = Required.Always)]
    public int Amount { get; set; }
    [JsonProperty("displayName", Required = Required.Always)]
    public string DisplayName { get; set; } = default!;
    [JsonProperty("tags", Required = Required.DisallowNull)]
    public string[]? Tags { get; set; } = default;
    [JsonProperty("name", Required = Required.Always)]
    public ModItemId ItemId { get; set; } = default!;
    public override string ToString() => Amount > 10000
        ? $"{Amount / 1000.0f:n2} B of {DisplayName}"
        : $"{Amount:n0} mB of {DisplayName}";
}

[JsonConverter(typeof(ModItemIdJsonConverter))]
[DebuggerDisplay($"{{{nameof(ToString)}(),nq}}")]
public class ModItemId {
    public ModItemId(string name) {
        var colon = name.IndexOf(':');
        if (colon < 0) throw new ArgumentException("Invalid mod item id!", nameof(name));
        ModName = name[..colon];
        ModItem = name[(colon + 1)..];
        if (ToString() != name) throw new InvalidProgramException("Bad Parsing!");
    }
    public override string ToString() => $"{ModName}:{ModItem}";
    public string ModName { get; }
    public string ModItem { get; }

    public class ModItemIdJsonConverter : JsonConverter<ModItemId> {
        public override ModItemId? ReadJson(JsonReader reader, Type objectType, ModItemId? existingValue, bool hasExistingValue, JsonSerializer serializer)
            => reader.Value is string value
                ? new(value)
                : throw new JsonException($"Could not parse mod name with token '{reader.Value}'");
        public override void WriteJson(JsonWriter writer, ModItemId? value, JsonSerializer serializer) {
            if (value is null) writer.WriteNull();
            else writer.WriteValue(value.ToString());
        }
    }
}

[JsonConverter(typeof(Md5JsonConverter))]
[DebuggerDisplay($"{{{nameof(ToString)}(),nq}}")]
public class Md5Hash : IEquatable<Md5Hash?> {
    private readonly byte[] _hash;
    public Md5Hash(string hash) : this(Convert.FromHexString(hash)) { }
    public Md5Hash(byte[] hash) {
        if (hash is not { Length: 16 }) throw new ArgumentException("Invalid digest size!", nameof(hash));
        _hash = hash;
    }
    public override bool Equals(object? obj) => Equals(obj as Md5Hash);
    public bool Equals(Md5Hash? other) => other != null && _hash.SequenceEqual(other._hash);
    public override int GetHashCode() {
        var hashCode = new HashCode();
        hashCode.AddBytes(_hash);
        return hashCode.ToHashCode();
    }
    public override string ToString() => Convert.ToHexString(_hash);

    public class Md5JsonConverter : JsonConverter<Md5Hash> {
        public override Md5Hash? ReadJson(JsonReader reader, Type objectType, Md5Hash? existingValue, bool hasExistingValue, JsonSerializer serializer)
            => reader.Value is string { Length: 32 } value
                ? new(value)
                : throw new JsonException($"Could not parse MD5 hash with token '{reader.Value}'");
        public override void WriteJson(JsonWriter writer, Md5Hash? value, JsonSerializer serializer) {
            if (value is null) writer.WriteNull();
            else writer.WriteValue(value.ToString());
        }
    }
}