Relay/RelayClient/Services/RelaySocketClient.cs

using System.Text.Json;
using RelayClient.Crypto;
using RelayShared.Services;
using WebSocketSharp;

namespace RelayClient.Services;

/// <summary>
/// The client-side WebSocket transport. Mirrors ChatSocketBehavior on the server.
///
/// Sending: typed helpers (SendGetHistory, SendRtcJoinChannel, SendEditMessage, …) build the
/// appropriate WsControlMessage or SocketEncryptedMessage and route through SendRaw. SendRaw
/// always uses synchronous _socket.Send because WebSocketSharp's SendAsync calls
/// Action.BeginInvoke internally, which throws PlatformNotSupportedException on .NET 5+.
/// Callers that need non-blocking sends (e.g. MainPage.SendMessage for image attachments)
/// wrap the call in Task.Run.
///
/// Receiving: OnMessage peeks the JSON. If it has an "Event" property → WsEventMessage (acks).
/// If it has a "Type" property → SignalType discriminator, deserialise into the right Socket*
/// type, fire the matching C# event. MainPage subscribes to these events.
///
/// Connect order matters: the first frame after the handshake is Authenticate (so the server
/// can verify the Core-issued token), then RegisterKey (so the server has our public key
/// before any encrypted message arrives), then GetServerKey + GetChannels.
/// </summary>
public sealed class RelaySocketClient
{
    /// <summary>Username this socket is authenticated as. Captured at construction.</summary>
    private readonly string _username;

    /// <summary>The underlying WebSocketSharp client. Owned (constructed) by this class.</summary>
    private readonly WebSocket _socket;

    /// <summary>
    /// The server's RSA public key, cached after the first GetServerKey response.
    /// MainPage reads this to encrypt outbound chat payloads.
    /// </summary>
    public string? ServerPublicKey { get; private set; }

    /// <summary>Fires for every raw incoming text frame. Mostly used for debug logging.</summary>
    public event Action<string>? RawMessageReceived;

    /// <summary>Fires when the server pushes a fresh channel list (initial connect or after CRUD).</summary>
    public event Action<SocketChannelList>? ChannelListReceived;

    /// <summary>Fires for newly-arrived chat messages (SignalType.EncryptedChat).</summary>
    public event Action<SocketEncryptedMessage>? EncryptedChatReceived;

    /// <summary>Fires when an existing message is edited by its author (SignalType.MessageEdited).</summary>
    public event Action<SocketEncryptedMessage>? MessageEdited;

    /// <summary>Fires when a message is deleted (SignalType.MessageDeleted).</summary>
    public event Action<SocketMessageDeletedEvent>? MessageDeleted;

    /// <summary>Fires when another user is typing in a channel.</summary>
    public event Action<SocketTypingEvent>? TypingReceived;

    /// <summary>Fires in response to a SendGetEditHistory request.</summary>
    public event Action<SocketEditHistoryResponse>? EditHistoryReceived;

    /// <summary>Fires for encrypted RTC SDP/ICE signals — RtcBridgeService forwards into the JS engine.</summary>
    public event Action<SocketRtcSignalMessage>? EncryptedRtcSignalReceived;

    /// <summary>Fires once when the server's public key arrives. Mainly used by tests; production reads ServerPublicKey directly.</summary>
    public event Action<string>? ServerPublicKeyReceived;

    /// <summary>Diagnostic logger. MainPage subscribes Console.WriteLine here.</summary>
    public event Action<string>? Log;

    /// <summary>Default URL points at localhost dev server. Production passes a remote URL.</summary>
    public RelaySocketClient(string username, string url = "ws://127.0.0.1:5001/")
    {
        _username = username;
        _socket = new WebSocket(url);
        _socket.OnMessage += OnMessage;
    }

    /// <summary>
    /// Opens the WebSocket and fires the four-step boot handshake IN ORDER:
    /// Authenticate → RegisterKey → GetServerKey → GetChannels. Order matters because the
    /// server uses RegisterKey to populate its session→username map (needed for permission
    /// checks on subsequent messages).
    /// </summary>
    public void Connect()
    {
        _socket.Connect();

        var publicKey = KeyStorage.LoadPublicKey(_username);

        SendControlMessage(new WsControlMessage { Action = WsAction.Authenticate, Username = _username, Token = MainPage._userToken });
        SendControlMessage(new WsControlMessage { Action = WsAction.RegisterKey,  Username = _username, PublicKey = publicKey });
        SendControlMessage(new WsControlMessage { Action = WsAction.GetServerKey });
        SendControlMessage(new WsControlMessage { Action = WsAction.GetChannels });
    }

    /// <summary>Detaches the message handler and closes the socket. Called from MainPage.OnDisappearing.</summary>
    public void Disconnect()
    {
        _socket.OnMessage -= OnMessage;
        if (_socket.ReadyState == WebSocketState.Open)
            _socket.Close();
    }

    /// <summary>Generic control-plane send. Serialises the WsControlMessage to JSON and ships it.</summary>
    public void SendControlMessage(WsControlMessage message) =>
        SendRaw(JsonSerializer.Serialize(message));

    /// <summary>Request the message history for a channel. Server streams it back as individual EncryptedChat frames.</summary>
    public void SendGetHistory(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.GetHistory, Username = _username, ChannelId = channelId });

    /// <summary>Tell the server we've joined a voice channel. Fires Speak permission check server-side.</summary>
    public void SendRtcJoinChannel(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.RtcJoin,  Username = _username, ChannelId = channelId });

    /// <summary>Tell the server we've left the voice channel. Idempotent server-side.</summary>
    public void SendRtcLeaveChannel(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.RtcLeave, Username = _username, ChannelId = channelId });

    /// <summary>Notify channel peers that we're typing. Server broadcasts a SocketTypingEvent to everyone but us.</summary>
    public void SendTyping(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.SendTyping, Username = _username, ChannelId = channelId });

    /// <summary>Request all historical versions of a message. Server replies with SocketEditHistoryResponse.</summary>
    public void SendGetEditHistory(string messageId, string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.GetEditHistory, Username = _username, MessageId = messageId, ChannelId = channelId });

    /// <summary>Create a new channel. Permission-gated server-side; on success the server broadcasts a fresh channel list.</summary>
    public void SendCreateChannel(string name, ChannelType type, string group = "") =>
        SendControlMessage(new WsControlMessage
        {
            Action = WsAction.CreateChannel,
            ChannelName = name,
            ChannelType = (int)type,
            ChannelGroup = group
        });

    /// <summary>Soft-delete a channel. Permission-gated server-side.</summary>
    public void SendDeleteChannel(string channelId) =>
        SendControlMessage(new WsControlMessage { Action = WsAction.DeleteChannel, ChannelId = channelId });

    /// <summary>
    /// Send an edit for an existing message. Caller is responsible for encrypting the new
    /// content (with the server's public key) before calling — same encryption shape as a new send.
    /// </summary>
    public void SendEditMessage(string messageId, string channelId, EncryptedPayload encrypted) =>
        SendJson(new SocketEncryptedMessage
        {
            Type = SignalType.ClientEditMessage, MessageId = messageId,
            SenderUsername = _username, ChannelId = channelId,
            CipherText = encrypted.CipherText, Nonce = encrypted.Nonce,
            Tag = encrypted.Tag, EncryptedKey = encrypted.EncryptedKey
        });

    /// <summary>Request soft-delete of one of our own messages. Server checks ownership before honoring.</summary>
    public void SendDeleteMessage(string messageId, string channelId) =>
        SendJson(new SocketEncryptedMessage
        {
            Type = SignalType.ClientDeleteMessage, MessageId = messageId,
            SenderUsername = _username, ChannelId = channelId
        });

    /// <summary>
    /// The single send pinch point. Synchronous (WebSocketSharp's SendAsync is broken on .NET 5+
    /// due to Action.BeginInvoke). All exceptions are logged AND rethrown so the calling
    /// Task.Run can surface them to the user via DisplayAlert.
    /// </summary>
    public void SendRaw(string message)
    {
        if (_socket.ReadyState != WebSocketState.Open)
        {
            Log?.Invoke($"[{_username}] Drop: socket not open ({_socket.ReadyState}), {message.Length} bytes.");
            return;
        }

        try
        {
            _socket.Send(message);
        }
        catch (Exception ex)
        {
            Log?.Invoke($"[{_username}] Send failed ({message.Length} bytes): {ex.Message}");
            throw;
        }
    }

    /// <summary>Convenience: JSON-serialise any payload and ship it. Used for all SocketEncryptedMessage and WsControlMessage sends.</summary>
    public void SendJson<T>(T payload) => SendRaw(JsonSerializer.Serialize(payload));

    /// <summary>
    /// WebSocketSharp callback for every incoming text frame. Peeks the JSON to decide whether
    /// it's a control-plane ack (Event property) or data-plane message (Type property), then
    /// fires the matching public C# event. Exceptions are caught locally so a malformed frame
    /// can't drop the connection.
    /// </summary>
    private void OnMessage(object? sender, MessageEventArgs e)
    {
        RawMessageReceived?.Invoke(e.Data);
        Log?.Invoke($"[{_username}] RAW: {e.Data[..Math.Min(200, e.Data.Length)]}");

        try
        {
            using var doc = JsonDocument.Parse(e.Data);
            var root = doc.RootElement;

            if (root.TryGetProperty("Event", out var evEl))
            {
                var wsEvent = (WsEvent)evEl.GetInt32();
                switch (wsEvent)
                {
                    case WsEvent.KeyRegistered: Log?.Invoke($"[{_username}] Key registered."); return;
                    case WsEvent.Authenticated: Log?.Invoke($"[{_username}] Authenticated.");  return;
                    case WsEvent.Error:
                        var det = root.TryGetProperty("Detail", out var d) ? d.GetString() : null;
                        Log?.Invoke($"[{_username}] Server error: {det}");
                        return;
                }
                return;
            }

            if (!root.TryGetProperty("Type", out var typeEl)) return;
            var type = (SignalType)typeEl.GetInt32();

            switch (type)
            {
                case SignalType.ChannelList:
                {
                    var p = JsonSerializer.Deserialize<SocketChannelList>(e.Data);
                    if (p is not null) ChannelListReceived?.Invoke(p);
                    return;
                }
                case SignalType.ServerPublicKey:
                {
                    var p = JsonSerializer.Deserialize<ServerPublicKeyMessage>(e.Data);
                    if (p is not null) { ServerPublicKey = p.PublicKey; ServerPublicKeyReceived?.Invoke(p.PublicKey); }
                    return;
                }
                case SignalType.EncryptedSignal:
                {
                    var p = JsonSerializer.Deserialize<SocketRtcSignalMessage>(e.Data);
                    if (p is not null) EncryptedRtcSignalReceived?.Invoke(p);
                    return;
                }
                case SignalType.EncryptedChat:
                {
                    var p = JsonSerializer.Deserialize<SocketEncryptedMessage>(e.Data);
                    if (p is not null) EncryptedChatReceived?.Invoke(p);
                    return;
                }
                case SignalType.MessageEdited:
                {
                    var p = JsonSerializer.Deserialize<SocketEncryptedMessage>(e.Data);
                    if (p is not null) MessageEdited?.Invoke(p);
                    return;
                }
                case SignalType.MessageDeleted:
                {
                    var p = JsonSerializer.Deserialize<SocketMessageDeletedEvent>(e.Data);
                    if (p is not null) MessageDeleted?.Invoke(p);
                    return;
                }
                case SignalType.TypingIndicator:
                {
                    var p = JsonSerializer.Deserialize<SocketTypingEvent>(e.Data);
                    if (p is not null) TypingReceived?.Invoke(p);
                    return;
                }
                case SignalType.EditHistory:
                {
                    var p = JsonSerializer.Deserialize<SocketEditHistoryResponse>(e.Data);
                    if (p is not null) EditHistoryReceived?.Invoke(p);
                    return;
                }
            }
        }
        catch (Exception ex)
        {
            Log?.Invoke($"[{_username}] WS parse error: {ex.Message}");
        }
    }
}