CellularManagement/src/CellularManagement.Infrastructure/Distributed/DistributedWebSocketManager.cs

using System.Net.WebSockets;
using CellularManagement.Domain.Entities;
using CellularManagement.Application.Services;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Caching.Distributed;
using System.Text.Json;

namespace CellularManagement.Infrastructure.Distributed;

public class DistributedWebSocketManager : IWebSocketService
{
    private readonly IDistributedCache _distributedCache;
    private readonly ILogger<DistributedWebSocketManager> _logger;
    private readonly string _nodeId;
    private const string CONNECTION_PREFIX = "ws_connection_";
    private const string USER_PREFIX = "ws_user_";
    private const string NODE_PREFIX = "ws_node_";

    public DistributedWebSocketManager(
        IDistributedCache distributedCache,
        ILogger<DistributedWebSocketManager> logger)
    {
        _distributedCache = distributedCache;
        _logger = logger;
        _nodeId = Guid.NewGuid().ToString();
    }

    public async Task<string> AcceptConnectionAsync(System.Net.WebSockets.WebSocket webSocket)
    {
        var connectionId = Guid.NewGuid().ToString();
        var connection = WebSocketConnection.Create(connectionId);
        
        var connectionKey = $"{CONNECTION_PREFIX}{connectionId}";
        var nodeKey = $"{NODE_PREFIX}{_nodeId}";
        
        // 存储连接信息
        await _distributedCache.SetStringAsync(
            connectionKey,
            JsonSerializer.Serialize(connection),
            new DistributedCacheEntryOptions
            {
                SlidingExpiration = TimeSpan.FromMinutes(30)
            });

        // 更新节点连接列表
        await AddConnectionToNodeAsync(connectionId);
        
        _logger.LogInformation("WebSocket connection accepted: {ConnectionId} on node {NodeId}", 
            connectionId, _nodeId);
        
        return connectionId;
    }

    public async Task<bool> CloseConnectionAsync(string connectionId)
    {
        var connectionKey = $"{CONNECTION_PREFIX}{connectionId}";
        var userKey = $"{USER_PREFIX}{connectionId}";
        
        // 获取连接信息
        var connectionJson = await _distributedCache.GetStringAsync(connectionKey);
        if (connectionJson != null)
        {
            var connection = JsonSerializer.Deserialize<WebSocketConnection>(connectionJson);
            if (connection != null)
            {
                connection.Close();
                await _distributedCache.SetStringAsync(
                    connectionKey,
                    JsonSerializer.Serialize(connection));
            }
        }

        // 从节点连接列表中移除
        await RemoveConnectionFromNodeAsync(connectionId);
        
        // 清理缓存
        await _distributedCache.RemoveAsync(connectionKey);
        await _distributedCache.RemoveAsync(userKey);
        
        _logger.LogInformation("WebSocket connection closed: {ConnectionId}", connectionId);
        return true;
    }

    public async Task<bool> SendMessageAsync(string connectionId, byte[] message)
    {
        var connectionKey = $"{CONNECTION_PREFIX}{connectionId}";
        var connectionJson = await _distributedCache.GetStringAsync(connectionKey);
        
        if (connectionJson != null)
        {
            var connection = JsonSerializer.Deserialize<WebSocketConnection>(connectionJson);
            if (connection?.State == WebSocketState.Open)
            {
                // 这里需要实现消息发送逻辑
                // 可能需要使用消息队列或其他机制
                return true;
            }
        }
        return false;
    }

    public async Task<bool> BroadcastMessageAsync(byte[] message)
    {
        // 获取所有节点的连接列表
        var nodes = await GetAllNodesAsync();
        foreach (var node in nodes)
        {
            // 向每个节点发送广播消息
            // 这里需要实现节点间的消息传递机制
        }
        return true;
    }

    public async Task<bool> SendMessageToUserAsync(string userId, byte[] message)
    {
        // 获取用户的所有连接
        var connections = await GetUserConnectionsAsync(userId);
        foreach (var connection in connections)
        {
            await SendMessageAsync(connection.ConnectionId, message);
        }
        return true;
    }

    public async Task AssociateUserAsync(string connectionId, string userId)
    {
        var connectionKey = $"{CONNECTION_PREFIX}{connectionId}";
        var userKey = $"{USER_PREFIX}{connectionId}";
        
        var connectionJson = await _distributedCache.GetStringAsync(connectionKey);
        if (connectionJson != null)
        {
            var connection = JsonSerializer.Deserialize<WebSocketConnection>(connectionJson);
            if (connection != null)
            {
                connection.AssociateUser(userId);
                await _distributedCache.SetStringAsync(
                    connectionKey,
                    JsonSerializer.Serialize(connection));
                await _distributedCache.SetStringAsync(userKey, userId);
            }
        }
    }

    public async Task<WebSocketConnection?> GetConnectionAsync(string connectionId)
    {
        var connectionKey = $"{CONNECTION_PREFIX}{connectionId}";
        var connectionJson = await _distributedCache.GetStringAsync(connectionKey);
        return connectionJson != null 
            ? JsonSerializer.Deserialize<WebSocketConnection>(connectionJson)
            : null;
    }

    public async Task<IEnumerable<WebSocketConnection>> GetUserConnectionsAsync(string userId)
    {
        var connections = new List<WebSocketConnection>();
        // 这里需要实现获取用户所有连接的逻辑
        return connections;
    }

    private async Task AddConnectionToNodeAsync(string connectionId)
    {
        var nodeKey = $"{NODE_PREFIX}{_nodeId}";
        var connections = await GetNodeConnectionsAsync();
        connections.Add(connectionId);
        await _distributedCache.SetStringAsync(
            nodeKey,
            JsonSerializer.Serialize(connections));
    }

    private async Task RemoveConnectionFromNodeAsync(string connectionId)
    {
        var nodeKey = $"{NODE_PREFIX}{_nodeId}";
        var connections = await GetNodeConnectionsAsync();
        connections.Remove(connectionId);
        await _distributedCache.SetStringAsync(
            nodeKey,
            JsonSerializer.Serialize(connections));
    }

    private async Task<List<string>> GetNodeConnectionsAsync()
    {
        var nodeKey = $"{NODE_PREFIX}{_nodeId}";
        var connectionsJson = await _distributedCache.GetStringAsync(nodeKey);
        return connectionsJson != null
            ? JsonSerializer.Deserialize<List<string>>(connectionsJson) ?? new List<string>()
            : new List<string>();
    }

    private async Task<List<string>> GetAllNodesAsync()
    {
        // 这里需要实现获取所有节点的逻辑
        // 可能需要使用服务发现或其他机制
        return new List<string>();
    }
}