AIGame/AIClient/Assets/HotUpdate/Scripts/Net/WsClient.cs

using System;
using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

namespace HotUpdate.Net
{
    public enum WsState
    {
        Disconnected,
        Connecting,
        Connected,
        Reconnecting,
    }

    /// <summary>
    /// 基于 ClientWebSocket 的传输（Editor / Standalone）。
    /// WebGL/微信需另接平台桥，本类在 WebGL 下不可用。
    /// </summary>
    public sealed class WsClient : IDisposable
    {
        readonly ConcurrentQueue<string> _inbox = new ConcurrentQueue<string>();
        readonly ConcurrentQueue<string> _outbox = new ConcurrentQueue<string>();

        ClientWebSocket _ws;
        CancellationTokenSource _cts;
        Task _recvTask;
        Task _sendTask;
        string _url;
        int _reconnectAttempt;
        bool _manualClose;

        public WsState State { get; private set; } = WsState.Disconnected;
        public string LastError { get; private set; }
        public event Action OnConnected;
        public event Action<string> OnDisconnected;
        public event Action<string> OnMessage;
        public event Action<int> OnReconnecting;

        public bool AutoReconnect { get; set; } = true;
        public float KeepAliveSeconds { get; set; } = 20f;

        public void Connect(string url)
        {
            _url = url;
            _manualClose = false;
            _reconnectAttempt = 0;
            _ = ConnectInternalAsync();
        }

        /// <summary>立刻断开并重新连接（重置退避计数）。</summary>
        public void ForceReconnect()
        {
            if (string.IsNullOrEmpty(_url)) return;
            _manualClose = false;
            AutoReconnect = true;
            _reconnectAttempt = 0;
            SafeCancel();
            try { _ws?.Dispose(); } catch { /* ignore */ }
            _ws = null;
            State = WsState.Disconnected;
            _ = ConnectInternalAsync();
        }

        public void SendText(string text)
        {
            if (string.IsNullOrEmpty(text)) return;
            _outbox.Enqueue(text);
        }

        public void Close()
        {
            _manualClose = true;
            AutoReconnect = false;
            SafeCancel();
            try { _ws?.Dispose(); } catch { /* ignore */ }
            _ws = null;
            State = WsState.Disconnected;
        }

        public void Tick()
        {
            while (_inbox.TryDequeue(out var msg))
                OnMessage?.Invoke(msg);
        }

        async Task ConnectInternalAsync()
        {
#if UNITY_WEBGL && !UNITY_EDITOR
            LastError = "ClientWebSocket not supported on WebGL; need WeChat bridge";
            State = WsState.Disconnected;
            OnDisconnected?.Invoke(LastError);
            return;
#else
            SafeCancel();
            _cts = new CancellationTokenSource();
            State = _reconnectAttempt > 0 ? WsState.Reconnecting : WsState.Connecting;
            LastError = null;

            try
            {
                _ws = new ClientWebSocket();
                _ws.Options.KeepAliveInterval = TimeSpan.FromSeconds(Mathf.Max(5f, KeepAliveSeconds));
                Debug.Log($"[WsClient] connecting {_url}");
                await _ws.ConnectAsync(new Uri(_url), _cts.Token);
                State = WsState.Connected;
                _reconnectAttempt = 0;
                Debug.Log("[WsClient] connected");
                OnConnected?.Invoke();

                _recvTask = RecvLoopAsync(_cts.Token);
                _sendTask = SendLoopAsync(_cts.Token);
                await Task.WhenAny(_recvTask, _sendTask);
            }
            catch (Exception e)
            {
                LastError = e.Message;
                Debug.LogWarning($"[WsClient] connect/loop error: {e.Message}");
            }
            finally
            {
                var wasConnected = State == WsState.Connected || State == WsState.Reconnecting || State == WsState.Connecting;
                State = WsState.Disconnected;
                try { _ws?.Dispose(); } catch { /* ignore */ }
                _ws = null;
                if (wasConnected)
                    OnDisconnected?.Invoke(LastError ?? "closed");

                if (!_manualClose && AutoReconnect)
                {
                    _reconnectAttempt++;
                    int delayMs = Math.Min(30000, 1000 * (1 << Math.Min(5, _reconnectAttempt)));
                    Debug.Log($"[WsClient] reconnect in {delayMs}ms attempt={_reconnectAttempt}");
                    OnReconnecting?.Invoke(_reconnectAttempt);
                    try
                    {
                        await Task.Delay(delayMs, CancellationToken.None);
                        if (!_manualClose && AutoReconnect)
                            _ = ConnectInternalAsync();
                    }
                    catch { /* ignore */ }
                }
            }
#endif
        }

        async Task RecvLoopAsync(CancellationToken ct)
        {
            var buffer = new byte[64 * 1024];
            var sb = new StringBuilder();
            while (!ct.IsCancellationRequested && _ws != null && _ws.State == WebSocketState.Open)
            {
                sb.Clear();
                WebSocketReceiveResult result;
                do
                {
                    result = await _ws.ReceiveAsync(new ArraySegment<byte>(buffer), ct);
                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        LastError = $"server close {_ws.CloseStatus} {_ws.CloseStatusDescription}";
                        return;
                    }
                    if (result.MessageType == WebSocketMessageType.Text)
                        sb.Append(Encoding.UTF8.GetString(buffer, 0, result.Count));
                } while (!result.EndOfMessage);

                if (sb.Length > 0)
                    _inbox.Enqueue(sb.ToString());
            }
        }

        async Task SendLoopAsync(CancellationToken ct)
        {
            while (!ct.IsCancellationRequested && _ws != null && _ws.State == WebSocketState.Open)
            {
                if (_outbox.TryDequeue(out var text))
                {
                    var bytes = Encoding.UTF8.GetBytes(text);
                    await _ws.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Text, true, ct);
                }
                else
                {
                    await Task.Delay(30, ct);
                }
            }
        }

        void SafeCancel()
        {
            try { _cts?.Cancel(); } catch { /* ignore */ }
            try { _cts?.Dispose(); } catch { /* ignore */ }
            _cts = null;
        }

        public void Dispose() => Close();
    }
}