Relay/RelayServer/Services/Chat/ChatSocketBehavior.cs

using System.Text.Json;
using RelayServer.Models;
using RelayServer.Services.Crypto;
using RelayServer.Services.Data;
using RelayServer.Services.Rtc;
using WebSocketSharp;
using WebSocketSharp.Server;
using ErrorEventArgs = WebSocketSharp.ErrorEventArgs;
using RelayShared.Services;

namespace RelayServer.Services.Chat;

/// <summary>
/// Handles websocket-based chat operations including client key registration,
/// server key retrieval, channel listing, channel history loading, and encrypted
/// channel message relay.
/// </summary>
public class ChatSocketBehavior : WebSocketBehavior
{
    public static ClientKeyService? ClientKeyService { get; set; }
    public static string? ServerPublicKey { get; set; }
    public static string? ServerPrivateKey { get; set; }
    public static string? ChannelDbKey { get; set; }
    public static ChannelCryptoService? ChannelCryptoService { get; set; }
    public static SurrealDb.Net.SurrealDbClient? Db { get; set; }

    /// <summary>
    /// Routes incoming websocket messages to the appropriate chat handler.
    /// </summary>
    /// <param name="e">The websocket message event arguments.</param>
    protected override void OnMessage(MessageEventArgs e)
    {
        var msg = e.Data;
        Console.WriteLine(msg);

        if (msg.StartsWith("REGISTER_KEY|"))
        {
            HandleRegisterKey(msg);
            return;
        }

        if (msg == "GET_SERVER_KEY")
        {
            HandleGetServerKey();
            return;
        }

        if (msg == "GET_CHANNELS")
        {
            HandleGetChannels();
            return;
        }

        if (msg.StartsWith("GET_HISTORY|"))
        {
            HandleGetHistory(msg);
            return;
        }
        
        if (msg.StartsWith("RTC_JOIN_CHANNEL|"))
        {
            HandleRtcJoinChannel(msg);
            return;
        }

        if (msg.StartsWith("RTC_LEAVE_CHANNEL|"))
        {
            HandleRtcLeaveChannel(msg);
            return;
        }

        if (IsEncryptedRtcSignal(msg))
        {
            HandleEncryptedRtcSignal(msg);
            return;
        }

        HandleEncryptedChatMessage(msg);
    }
    
    private static bool IsEncryptedRtcSignal(string msg)
    {
        try
        {
            using var doc = JsonDocument.Parse(msg);
            var root = doc.RootElement;

            if (!root.TryGetProperty("Type", out var typeProp))
                return false;

            var type = (SignalType)typeProp.GetInt32();

            return type == SignalType.EncryptedSignal;
        }
        catch
        {
            return false;
        }
    }
    
    private void HandleEncryptedRtcSignal(string msg)
    {
        Console.WriteLine("RTC SIGNAL HIT");
        SocketRtcSignalMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketRtcSignalMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted RTC signal payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != SignalType.EncryptedSignal)
            return;

        if (string.IsNullOrWhiteSpace(clientPayload.ChannelId))
        {
            Console.WriteLine("Encrypted RTC signal missing channel id.");
            return;
        }

        string plainText;

        try
        {
            plainText = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt RTC signal: {ex.Message}");
            return;
        }

        var sessionIds = RtcChannelPresenceService.GetSessionsInChannel(clientPayload.ChannelId);

        foreach (var sessionId in sessionIds)
        {
            if (sessionId == ID)
                continue;

            var username = RtcChannelPresenceService.GetUsernameForSession(sessionId);
            if (string.IsNullOrWhiteSpace(username))
                continue;

            var clientKey = GetClientPublicKeyByUsernameSync(username);
            if (clientKey is null)
                continue;

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, clientKey.PublicKey);

            var outbound = new SocketRtcSignalMessage
            {
                Type = SignalType.EncryptedSignal,
                SenderUsername = clientPayload.SenderUsername,
                ChannelId = clientPayload.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Sessions.SendTo(JsonSerializer.Serialize(outbound), sessionId);
        }

        Console.WriteLine($"Forwarded encrypted RTC signal from {clientPayload.SenderUsername} to channel {clientPayload.ChannelId}");
    }
    
    /// <summary>
    /// 
    /// </summary>
    /// <param name="e"></param>
    protected override void OnClose(CloseEventArgs e)
    {
        RtcChannelPresenceService.RemoveSession(ID);
        Console.WriteLine($"WebSocket closed: session={ID}, code={e.Code}, reason={e.Reason}");
        base.OnClose(e);
    }
    
    protected override void OnError(ErrorEventArgs e)
    {
        Console.WriteLine($"WebSocket error: session={ID}, message={e.Message}");
        base.OnError(e);
    }

    /// <summary>
    /// Extracts a display username from a stored user record id value.
    /// </summary>
    /// <param name="senderUserId">The stored sender user id.</param>
    /// <returns>
    /// The extracted username when possible; otherwise, a fallback value.
    /// </returns>
    private static string ExtractUsernameFromUserId(string senderUserId)
    {
        if (string.IsNullOrWhiteSpace(senderUserId))
            return "Unknown";

        var parts = senderUserId.Split(':', 2);
        return parts.Length == 2 ? parts[1] : senderUserId;
    }

    /// <summary>
    /// Registers or updates a client's public key from a websocket registration payload.
    /// </summary>
    /// <param name="msg">The raw websocket registration message.</param>
    private void HandleRegisterKey(string msg)
    {
        var parts = msg.Split('|', 3);

        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid REGISTER_KEY payload.");
            return;
        }

        var username = parts[1];
        var publicKey = parts[2];

        if (ClientKeyService is null)
        {
            Console.WriteLine("ClientKeyService is not initialized.");
            return;
        }

        RegisterOrUpdateClientKeySync(username, publicKey);

        Send($"SERVER:REGISTERED_KEY:{username}");
    }

    /// <summary>
    /// Sends the current list of channels to the connected websocket client.
    /// </summary>
    private void HandleGetChannels()
    {
        if (Db is null)
        {
            Console.WriteLine("Db is not initialized.");
            return;
        }
        //TODO: Update to include ChannelType and Group String on channels
        var channels = GetChannelsSync()
            .OrderBy(c => c.CreatedAt)
            .Select(c => new ChannelItem()
            {
                ChannelId = GetRecordId(c.Id),
                Name = c.Name,
                CreatedAt = c.CreatedAt
            })
            .ToList();

        var payload = new SocketChannelList
        {
            Type = SignalType.ChannelList,
            Channels = channels
        };

        Send(JsonSerializer.Serialize(payload));
    }

    /// <summary>
    /// Sends the server's public key to the connected websocket client.
    /// </summary>
    private void HandleGetServerKey()
    {
        if (string.IsNullOrWhiteSpace(ServerPublicKey))
        {
            Console.WriteLine("Server public key is not initialized.");
            return;
        }

        var payload = new ServerPublicKeyMessage
        {
            Type = SignalType.ServerPublicKey,
            PublicKey = ServerPublicKey
        };

        Send(JsonSerializer.Serialize(payload));
    }

    /// <summary>
    /// Decrypts an incoming encrypted chat payload, stores it in the database,
    /// and rebroadcasts it to connected clients encrypted with each client's public key.
    /// </summary>
    /// <param name="msg">The raw encrypted chat websocket message.</param>
    private void HandleEncryptedChatMessage(string msg)
    {
        SocketEncryptedMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketEncryptedMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted client payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != SignalType.ClientEncryptedChat)
            return;

        if (!EnsureCoreReady() || !EnsureCryptoReady())
            return;

        string plainText;

        try
        {
            plainText = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt client payload: {ex.Message}");
            return;
        }

        Console.WriteLine($"Server decrypted message from {clientPayload.SenderUsername}: {plainText}");
        try
        {
            var dbEncrypted = ChannelCryptoService.Encrypt(plainText, ChannelDbKey);

            var savedMessage = CreateChannelMessageSync(new ChannelMessages
            {
                ChannelId = clientPayload.ChannelId,
                SenderUserId = $"users:{clientPayload.SenderUsername.ToLower()}",
                CipherText = dbEncrypted.cipherText,
                Nonce = dbEncrypted.nonce,
                Tag = dbEncrypted.tag,
                CreatedAt = DateTime.UtcNow
            });

            Console.WriteLine($"Live message saved to DB: {JsonSerializer.Serialize(savedMessage)}");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to save live message to DB: {ex.Message}");
            return;
        }

        var allKeys = GetAllClientPublicKeysSync();

        foreach (var client in allKeys)
        {
            var encrypted = E2EeHelper.EncryptForRecipient(plainText, client.PublicKey);

            Console.WriteLine($"Encrypting outbound message from {clientPayload.SenderUsername} for {client.Username}");

            var outbound = new SocketEncryptedMessage
            {
                Type = SignalType.EncryptedChat,
                SenderUsername = clientPayload.SenderUsername,
                RecipientUsername = client.Username,
                ChannelId = clientPayload.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Sessions.Broadcast(JsonSerializer.Serialize(outbound));
        }
    }

    /// <summary>
    /// Loads stored channel history for a specific user and channel, decrypts it from
    /// database storage format, and sends it back encrypted for the requesting client.
    /// </summary>
    /// <param name="msg">The raw history request websocket message.</param>
    private void HandleGetHistory(string msg)
    {
        var parts = msg.Split('|', 3);

        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid GET_HISTORY payload.");
            return;
        }

        var username = parts[1];
        var channelId = parts[2];

        if (!EnsureCoreReady() || ChannelCryptoService is null || string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("History dependencies are not initialized.");
            return;
        }

        var targetClient = GetClientPublicKeyByUsernameSync(username);

        if (targetClient is null)
        {
            Console.WriteLine($"No public key found for history request user {username}");
            return;
        }

        var allMessages = GetChannelMessagesSync();

        var channelMessages = allMessages
            .Where(m => m.ChannelId == channelId)
            .OrderBy(m => m.CreatedAt)
            .ToList();

        Console.WriteLine($"Sending {channelMessages.Count} history messages to {username}");

        foreach (var dbMessage in channelMessages)
        {
            string plainText;

            try
            {
                plainText = ChannelCryptoService.Decrypt(
                    dbMessage.CipherText,
                    dbMessage.Nonce,
                    dbMessage.Tag,
                    ChannelDbKey
                );
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Failed to decrypt DB history row {dbMessage.Id}: {ex.Message}");
                continue;
            }

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, targetClient.PublicKey);

            var outbound = new SocketEncryptedMessage
            {
                Type = SignalType.EncryptedChat,
                SenderUsername = ExtractUsernameFromUserId(dbMessage.SenderUserId),
                RecipientUsername = username,
                ChannelId = channelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Send(JsonSerializer.Serialize(outbound));
        }
    }

    /// <summary>
    /// Converts a SurrealDB record id object into a table:id string representation.
    /// </summary>
    /// <param name="id">The raw record id object.</param>
    /// <returns>
    /// A formatted record id string, or an empty string if the input is null.
    /// </returns>
    private static string GetRecordId(object? id)
    {
        if (id is null)
            return string.Empty;

        var json = JsonSerializer.Serialize(id);

        using var doc = JsonDocument.Parse(json);

        var root = doc.RootElement;

        var recordId = root.GetProperty("Id").GetString() ?? string.Empty;
        var table = root.GetProperty("Table").GetString() ?? string.Empty;

        return $"{table}:{recordId}";
    }
    
    /// <summary>
    /// Synchronously registers or updates a stored client public key using the async key service.
    /// </summary>
    /// <param name="username">The client username.</param>
    /// <param name="publicKey">The client's public key.</param>
    private void RegisterOrUpdateClientKeySync(string username, string publicKey)
    {
        Task.Run(async () => await ClientKeyService!.RegisterOrUpdateKeyAsync(username, publicKey))
            .GetAwaiter()
            .GetResult();
    }

    /// <summary>
    /// Synchronously loads all channels from the database.
    /// </summary>
    /// <returns>A list of channel records.</returns>
    private List<Channels> GetChannelsSync()
    {
        return Task.Run(async () => await Db!.Select<Channels>("channels"))
            .GetAwaiter()
            .GetResult()
            .ToList();
    }

    /// <summary>
    /// Synchronously gets the stored public key record for the specified user.
    /// </summary>
    /// <param name="username">The username to look up.</param>
    /// <returns>
    /// The matching client public key record, or null if none exists.
    /// </returns>
    private ClientPublicKeys? GetClientPublicKeyByUsernameSync(string username)
    {
        return Task.Run(async () => await ClientKeyService!.GetByUsernameAsync(username))
            .GetAwaiter()
            .GetResult();
    }

    /// <summary>
    /// Synchronously loads all stored client public key records.
    /// </summary>
    /// <returns>A list of all client public key records.</returns>
    private List<ClientPublicKeys> GetAllClientPublicKeysSync()
    {
        return Task.Run(async () => await ClientKeyService!.GetAllAsync())
            .GetAwaiter()
            .GetResult();
    }

    /// <summary>
    /// Synchronously loads all stored channel messages from the database.
    /// </summary>
    /// <returns>A list of channel message records.</returns>
    private List<ChannelMessages> GetChannelMessagesSync()
    {
        return Task.Run(async () => await Db!.Select<ChannelMessages>("channel_messages"))
            .GetAwaiter()
            .GetResult()
            .ToList();
    }

    /// <summary>
    /// Synchronously creates a new channel message record in the database.
    /// </summary>
    /// <param name="message">The message record to create.</param>
    /// <returns>The created channel message record.</returns>
    private ChannelMessages CreateChannelMessageSync(ChannelMessages message)
    {
        return Task.Run(async () => await Db!.Create("channel_messages", message))
            .GetAwaiter()
            .GetResult();
    }
    
    /// <summary>
    /// 
    /// </summary>
    /// <returns></returns>
    private bool EnsureCoreReady()
    {
        if (ClientKeyService is null || Db is null)
        {
            Console.WriteLine("Core services not initialized.");
            return false;
        }

        return true;
    }
    
    /// <summary>
    /// 
    /// </summary>
    /// <returns></returns>
    private bool EnsureCryptoReady()
    {
        if (string.IsNullOrWhiteSpace(ServerPrivateKey) || string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("Crypto keys not initialized.");
            return false;
        }

        if (ChannelCryptoService is null)
        {
            Console.WriteLine("ChannelCryptoService is not initialized.");
            return false;
        }

        return true;
    }
    
    /// <summary>
    /// 
    /// </summary>
    /// <param name="msg"></param>
    private void HandleRtcJoinChannel(string msg)
    {
        var parts = msg.Split('|', 3);
        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid RTC_JOIN_CHANNEL payload.");
            return;
        }

        var username = parts[1];
        var channelId = parts[2];

        RtcChannelPresenceService.SetUser(ID, username);
        RtcChannelPresenceService.JoinChannel(ID, channelId);

        Console.WriteLine($"RTC presence joined: session={ID}, user={username}, channel={channelId}");
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="msg"></param>
    private void HandleRtcLeaveChannel(string msg)
    {
        var parts = msg.Split('|', 3);
        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid RTC_LEAVE_CHANNEL payload.");
            return;
        }

        var username = parts[1];
        var channelId = parts[2];

        if (RtcChannelPresenceService.IsInChannel(ID, channelId))
        {
            RtcChannelPresenceService.LeaveChannel(ID);
        }

        Console.WriteLine($"RTC presence left: session={ID}, user={username}, channel={channelId}");
    }
}