maki/Maki/Gateway/GatewayShard.cs

using Maki.Structures.Channels;
using Maki.Structures.Gateway;
using Maki.Structures.Guilds;
using Maki.Structures.Messages;
using Maki.Structures.Presences;
using Maki.Structures.Users;
using Newtonsoft.Json;
using System;
using WebSocketSharp;

namespace Maki.Gateway
{
    /// <summary>
    /// Gateway connection shard
    /// </summary>
    class GatewayShard : IDisposable
    {
        /// <summary>
        /// Session key for continuing a resuming after disconnecting
        /// </summary>
        private string session;

        /// <summary>
        /// Websocket container
        /// </summary>
        private WebSocket WebSocket;
        
        /// <summary>
        /// Interval at which heartbeats are sent
        /// </summary>
        public TimeSpan HeartbeatInterval { get; private set; }

        /// <summary>
        /// Last sequence received from the gateway, primarily used for resuming connections and getting the data that was missed
        /// </summary>
        public int? LastSequence { get; private set; } = null;

        /// <summary>
        /// Identifier for this Shard
        /// </summary>
        public readonly int Id;

        /// <summary>
        /// Parent DiscordClient instance
        /// </summary>
        private readonly Discord client;

        /// <summary>
        /// Heartbeat handler
        /// </summary>
        private readonly GatewayHeartbeatManager heartbeatHandler;

        #region Events
        public event Action<GatewayShard> OnCallCreate;
        public event Action<GatewayShard> OnCallDelete;
        public event Action<GatewayShard> OnCallUpdate;
        public event Action<GatewayShard> OnVoiceServerUpdate;
        public event Action<GatewayShard> OnVoiceStateUpdate;

        public event Action<GatewayShard, Channel> OnChannelCreate;
        public event Action<GatewayShard, Channel> OnChannelDelete;
        public event Action<GatewayShard> OnChannelPinsAck;
        public event Action<GatewayShard> OnChannelPinsUpdate;
        public event Action<GatewayShard> OnChannelRecipientAdd;
        public event Action<GatewayShard> OnChannelRecipientRemove;
        public event Action<GatewayShard, Channel> OnChannelUpdate;

        public event Action<GatewayShard, User> OnGuildBanAdd;
        public event Action<GatewayShard, User> OnGuildBanRemove;
        public event Action<GatewayShard, Guild> OnGuildCreate;
        public event Action<GatewayShard, Guild> OnGuildDelete;
        public event Action<GatewayShard, Guild> OnGuildEmojisUpdate;
        public event Action<GatewayShard, GuildIntegration> OnGuildIntegrationsUpdate;
        public event Action<GatewayShard, GuildMember> OnGuildMemberAdd;
        public event Action<GatewayShard, GuildMember> OnGuildMemberUpdate;
        public event Action<GatewayShard, GuildMember> OnGuildMemberRemove;
        public event Action<GatewayShard, GuildMembersChunk> OnGuildMembersChunk;
        public event Action<GatewayShard, GuildRole> OnGuildRoleCreate;
        public event Action<GatewayShard, GuildRole> OnGuildRoleUpdate;
        public event Action<GatewayShard, GuildRole> OnGuildRoleDelete;
        public event Action<GatewayShard, Guild> OnGuildUpdate;

        public event Action<GatewayShard> OnMessageAck;
        public event Action<GatewayShard, Message> OnMessageCreate;
        public event Action<GatewayShard, Message> OnMessageDelete;
        public event Action<GatewayShard> OnMessageDeleteBulk;
        public event Action<GatewayShard> OnMessageReactionAdd;
        public event Action<GatewayShard> OnMessageReactionRemove;
        public event Action<GatewayShard> OnMessageReactionsRemoveAll;
        public event Action<GatewayShard, Message> OnMessageUpdate;

        public event Action<GatewayShard, TypingStart> OnTypingStart;
        public event Action<GatewayShard, Presence> OnPresenceUpdate;
        public event Action<GatewayShard> OnPresencesReplace;

        public event Action<GatewayShard> OnFriendSuggestionDelete;
        public event Action<GatewayShard> OnRelationshipAdd;
        public event Action<GatewayShard> OnRelationshipRemove;

        public event Action<GatewayShard, GatewayReady> OnReady;
        public event Action<GatewayShard> OnResumed;

        public event Action<GatewayShard, User> OnUserUpdate;
        public event Action<GatewayShard> OnUserNoteUpdate;
        public event Action<GatewayShard> OnUserSettingsUpdate;

        public event Action<GatewayShard> OnSocketOpen;
        public event Action<GatewayShard, bool, ushort, string> OnSocketClose;
        public event Action<GatewayShard, Exception> OnSocketError;
        public event Action<GatewayShard, string> OnSocketMessage;
        #endregion

        /// <summary>
        /// Constructor
        /// </summary>
        /// <param name="id">Shard Id</param>
        /// <param name="c">Parent DiscordClient instance</param>
        public GatewayShard(int id, Discord c)
        {
            Id = id;
            client = c;
            heartbeatHandler = new GatewayHeartbeatManager(this);
            Connect();
        }

        /// <summary>
        /// Event handler for WebSocketSharp's OnOpen event, forwards to our OnSocketOpen event
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnOpen(object sender, EventArgs e) => OnSocketOpen?.Invoke(this);

        /// <summary>
        /// Event handler for WebSocketSharp's OnError event, forwards to our OnSocketError event
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnError(object sender, ErrorEventArgs e) => OnSocketError?.Invoke(this, e.Exception);

        /// <summary>
        /// Event handler for WebSocketSharp's OnClose event, stops heartbeats, forwards to our OnSocketClose event and reconnects if the close wasn't clean
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnClose(object sender, CloseEventArgs e)
        {
            heartbeatHandler.Stop();
            OnSocketClose?.Invoke(this, e.WasClean, e.Code, e.Reason);

            if (!e.WasClean || e.Code != 1000)
                Connect();
        }

        /// <summary>
        /// Event handler for WebSocketSharp's OnOpen event, handles gateway instructions and forwards to our events
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnMessage(object sender, MessageEventArgs e)
        {
            // ignore non-text messages
            if (!e.IsText)
                return;

            Console.WriteLine(e.Data.Replace(client.Token, new string('*', client.Token.Length)));
            OnSocketMessage?.Invoke(this, e.Data);
            GatewayPayload payload = JsonConvert.DeserializeObject<GatewayPayload>(e.Data);
            
            switch (payload.OPCode) {
                case GatewayOPCode.Dispatch:
                    LastSequence = payload.Sequence;
                    Enum.TryParse(payload.Name, out GatewayEvent evt);

                    switch (evt) {
                        #region Call
                        case GatewayEvent.CALL_CREATE:
                            OnCallCreate?.Invoke(this);
                            break;

                        case GatewayEvent.CALL_DELETE:
                            OnCallDelete?.Invoke(this);
                            break;

                        case GatewayEvent.CALL_UPDATE:
                            OnCallUpdate?.Invoke(this);
                            break;

                        case GatewayEvent.VOICE_SERVER_UPDATE:
                            OnVoiceServerUpdate?.Invoke(this);
                            break;

                        case GatewayEvent.VOICE_STATE_UPDATE:
                            OnVoiceStateUpdate?.Invoke(this);
                            break;
                        #endregion

                        #region Channel
                        case GatewayEvent.CHANNEL_CREATE:
                            OnChannelCreate?.Invoke(this, payload.DataAs<Channel>());
                            break;

                        case GatewayEvent.CHANNEL_DELETE:
                            OnChannelDelete?.Invoke(this, payload.DataAs<Channel>());
                            break;

                        case GatewayEvent.CHANNEL_PINS_ACK:
                            OnChannelPinsAck?.Invoke(this);
                            break;

                        case GatewayEvent.CHANNEL_PINS_UPDATE:
                            OnChannelPinsUpdate?.Invoke(this);
                            break;

                        case GatewayEvent.CHANNEL_RECIPIENT_ADD:
                            OnChannelRecipientAdd?.Invoke(this);
                            break;

                        case GatewayEvent.CHANNEL_RECIPIENT_REMOVE:
                            OnChannelRecipientRemove?.Invoke(this);
                            break;

                        case GatewayEvent.CHANNEL_UPDATE:
                            OnChannelUpdate?.Invoke(this, payload.DataAs<Channel>());
                            break;
                        #endregion

                        #region Guild
                        case GatewayEvent.GUILD_BAN_ADD:
                            OnGuildBanAdd?.Invoke(this, payload.DataAs<User>());
                            break;

                        case GatewayEvent.GUILD_BAN_REMOVE:
                            OnGuildBanRemove?.Invoke(this, payload.DataAs<User>());
                            break;

                        case GatewayEvent.GUILD_CREATE:
                            OnGuildCreate?.Invoke(this, payload.DataAs<Guild>());
                            break;

                        case GatewayEvent.GUILD_DELETE:
                            OnGuildDelete?.Invoke(this, payload.DataAs<Guild>());
                            break;

                        case GatewayEvent.GUILD_EMOJIS_UPDATE:
                            OnGuildEmojisUpdate?.Invoke(this, payload.DataAs<Guild>());
                            break;

                        case GatewayEvent.GUILD_INTEGRATIONS_UPDATE:
                            OnGuildIntegrationsUpdate?.Invoke(this, payload.DataAs<GuildIntegration>());
                            break;

                        case GatewayEvent.GUILD_MEMBER_ADD:
                            OnGuildMemberAdd?.Invoke(this, payload.DataAs<GuildMember>());
                            break;

                        case GatewayEvent.GUILD_MEMBER_UPDATE:
                            OnGuildMemberUpdate?.Invoke(this, payload.DataAs<GuildMember>());
                            break;

                        case GatewayEvent.GUILD_MEMBER_REMOVE:
                            OnGuildMemberRemove?.Invoke(this, payload.DataAs<GuildMember>());
                            break;

                        case GatewayEvent.GUILD_MEMBERS_CHUNK:
                            OnGuildMembersChunk?.Invoke(this, payload.DataAs<GuildMembersChunk>());
                            break;

                        case GatewayEvent.GUILD_ROLE_CREATE:
                            OnGuildRoleCreate?.Invoke(this, payload.DataAs<GuildRole>());
                            break;

                        case GatewayEvent.GUILD_ROLE_UPDATE:
                            OnGuildRoleUpdate?.Invoke(this, payload.DataAs<GuildRole>());
                            break;

                        case GatewayEvent.GUILD_ROLE_DELETE:
                            OnGuildRoleDelete?.Invoke(this, payload.DataAs<GuildRole>());
                            break;

                        case GatewayEvent.GUILD_UPDATE:
                            OnGuildUpdate?.Invoke(this, payload.DataAs<Guild>());
                            break;
                        #endregion

                        #region Message
                        case GatewayEvent.MESSAGE_ACK:
                            OnMessageAck?.Invoke(this);
                            break;

                        case GatewayEvent.MESSAGE_CREATE:
                            OnMessageCreate?.Invoke(this, payload.DataAs<Message>());
                            break;

                        case GatewayEvent.MESSAGE_DELETE:
                            OnMessageDelete?.Invoke(this, payload.DataAs<Message>());
                            break;

                        case GatewayEvent.MESSAGE_DELETE_BULK:
                            OnMessageDeleteBulk?.Invoke(this);
                            break;

                        case GatewayEvent.MESSAGE_REACTION_ADD:
                            OnMessageReactionAdd?.Invoke(this);
                            break;

                        case GatewayEvent.MESSAGE_REACTION_REMOVE:
                            OnMessageReactionRemove?.Invoke(this);
                            break;

                        case GatewayEvent.MESSAGE_REACTIONS_REMOVE_ALL:
                            OnMessageReactionsRemoveAll?.Invoke(this);
                            break;

                        case GatewayEvent.MESSAGE_UPDATE:
                            OnMessageUpdate?.Invoke(this, payload.DataAs<Message>());
                            break;
                        #endregion

                        #region Presence
                        case GatewayEvent.TYPING_START:
                            OnTypingStart?.Invoke(this, payload.DataAs<TypingStart>());
                            break;

                        case GatewayEvent.PRESENCE_UPDATE:
                            OnPresenceUpdate?.Invoke(this, payload.DataAs<Presence>());
                            break;

                        case GatewayEvent.PRESENCES_REPLACE:
                            OnPresencesReplace?.Invoke(this);
                            break;
                        #endregion

                        #region Relations
                        case GatewayEvent.FRIEND_SUGGESTION_DELETE:
                            OnFriendSuggestionDelete?.Invoke(this);
                            break;

                        case GatewayEvent.RELATIONSHIP_ADD:
                            OnRelationshipAdd?.Invoke(this);
                            break;

                        case GatewayEvent.RELATIONSHIP_REMOVE:
                            OnRelationshipRemove?.Invoke(this);
                            break;
                        #endregion

                        #region State
                        case GatewayEvent.READY:
                            GatewayReady ready = payload.DataAs<GatewayReady>();
                            session = ready.Session;
                            OnReady?.Invoke(this, ready);
                            break;

                        case GatewayEvent.RESUMED:
                            OnResumed?.Invoke(this);
                            break;
                        #endregion

                        #region User
                        case GatewayEvent.USER_UPDATE:
                            OnUserUpdate?.Invoke(this, payload.DataAs<User>());
                            break;

                        case GatewayEvent.USER_NOTE_UPDATE:
                            OnUserNoteUpdate?.Invoke(this);
                            break;

                        case GatewayEvent.USER_SETTINGS_UPDATE:
                            OnUserSettingsUpdate?.Invoke(this);
                            break;
                        #endregion

                        default:
                            Console.WriteLine($"Unknown payload type: {payload.Name}");
                            break;
                    }
                    break;

                case GatewayOPCode.Hello:
                    GatewayHello hello = payload.DataAs<GatewayHello>();

                    HeartbeatInterval = TimeSpan.FromMilliseconds(hello.HeartbeatInterval);
                    heartbeatHandler.Start();

                    if (string.IsNullOrEmpty(session))
                        Send(GatewayOPCode.Identify, new GatewayIdentification {
                            Token = client.Token,
                            Compress = false,
                            LargeThreshold = 250,
                            Shard = new int[] { Id, client.ShardClient.ShardCount },
                            Properties = new GatewayIdentificationProperties {
                                OperatingSystem = "windows",
                                Browser = "Maki",
                                Device = "Maki",
                                Referrer = string.Empty,
                                ReferringDomain = string.Empty,
                            }
                        });
                    else
                        Send(GatewayOPCode.Resume, new GatewayResume {
                            LastSequence = LastSequence ?? default(int),
                            Token = client.Token,
                            Session = session,
                        });
                    break;
            }
        }

        #region IDisposable

        private bool IsDisposed = false;

        private void Dispose(bool disposing)
        {
            if (IsDisposed)
                return;

            IsDisposed = true;
            Disconnect();

            if (disposing)
                GC.SuppressFinalize(this);
        }

        ~GatewayShard()
            => Dispose(false);

        public void Dispose()
            => Dispose(true);

        #endregion

        /// <summary>
        /// Connects to the gateway
        /// </summary>
        public void Connect()
        {
            WebSocket = new WebSocket($"{client.Gateway}?v={Discord.GATEWAY_VERSION}&encoding=json");

            // make wss not log anything on its own
            WebSocket.Log.Output = (LogData logData, string path) => { };

            WebSocket.OnOpen += WebSocket_OnOpen;
            WebSocket.OnClose += WebSocket_OnClose;
            WebSocket.OnError += WebSocket_OnError;
            WebSocket.OnMessage += WebSocket_OnMessage;
        
            WebSocket.Connect();
        }

        /// <summary>
        /// Stops heartbeats and closes gateway connection for this shard
        /// </summary>
        public void Disconnect()
        {
            heartbeatHandler.Stop();

            if (WebSocket.ReadyState != WebSocketState.Closed)
                WebSocket?.Close(CloseStatusCode.Normal);
        }

        /// <summary>
        /// Serialises and sends a json object
        /// </summary>
        /// <typeparam name="T">Type to serialise as</typeparam>
        /// <param name="data">Data to serialise and send</param>
        public void Send<T>(T data)
        {
            string json = JsonConvert.SerializeObject(data, typeof(T), new JsonSerializerSettings());

            Console.WriteLine(json.Replace(client.Token, new string('*', client.Token.Length)));

            WebSocket.Send(json);
        }

        /// <summary>
        /// Serialises and sends a Gateway Payload
        /// </summary>
        /// <param name="opcode">Opcode to use</param>
        /// <param name="data">Data to send</param>
        public void Send(GatewayOPCode opcode, object data)
        {
            Send(new GatewayPayload
            {
                OPCode = opcode,
                Data = data
            });
        }
    }
}