using System;
using System.Collections.Concurrent;
#if !UNITY_WEBGL || UNITY_EDITOR
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
#endif
using AOT.Boot;
using UnityEngine;

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

    /// <summary>
    /// Editor/Standalone 使用 ClientWebSocket；WebGL 使用微信 SocketTask/浏览器桥。
    /// </summary>
    public sealed class WsClient : IDisposable
    {
        readonly ConcurrentQueue<string> _inbox = new ConcurrentQueue<string>();
        readonly ConcurrentQueue<string> _outbox = new ConcurrentQueue<string>();

#if UNITY_WEBGL && !UNITY_EDITOR
        int _browserSocket;
        float _connectDeadline;
        float _reconnectAt = -1;
#else
        ClientWebSocket _ws;
        CancellationTokenSource _cts;
        Task _recvTask;
        Task _sendTask;
#endif
        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;
#if UNITY_WEBGL && !UNITY_EDITOR
            StartBrowserConnection();
#else
            _ = ConnectInternalAsync();
#endif
        }

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

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

        public void Close()
        {
            _manualClose = true;
            AutoReconnect = false;
#if UNITY_WEBGL && !UNITY_EDITOR
            ReleaseBrowserSocket();
            _reconnectAt = -1;
#else
            SafeCancel();
            try { _ws?.Dispose(); } catch { /* ignore */ }
            _ws = null;
#endif
            State = WsState.Disconnected;
        }

        public void Tick()
        {
#if UNITY_WEBGL && !UNITY_EDITOR
            TickBrowser();
#endif
            while (_inbox.TryDequeue(out var msg))
                OnMessage?.Invoke(msg);
        }

#if UNITY_WEBGL && !UNITY_EDITOR
        void ReleaseBrowserSocket()
        {
            if (_browserSocket != 0) WebSocketBridge.AIGame_WsDispose(_browserSocket);
            _browserSocket = 0;
        }

        void StartBrowserConnection()
        {
            ReleaseBrowserSocket();
            _reconnectAt = -1;
            LastError = null;
            State = _reconnectAttempt > 0 ? WsState.Reconnecting : WsState.Connecting;
            _connectDeadline = Time.realtimeSinceStartup + 15f;
            _browserSocket = WebSocketBridge.AIGame_WsCreate(_url);
            if (_browserSocket == 0) BrowserDisconnected("WebSocket bridge unavailable");
        }

        void BrowserDisconnected(string reason)
        {
            ReleaseBrowserSocket();
            LastError = reason;
            State = WsState.Disconnected;
            while (_outbox.TryDequeue(out _)) { }
            _reconnectAt = -1;
            if (!_manualClose && AutoReconnect)
            {
                _reconnectAttempt++;
                _reconnectAt = Time.realtimeSinceStartup + Math.Min(30, 1 << Math.Min(5, _reconnectAttempt));
            }
            OnDisconnected?.Invoke(reason);
            if (_reconnectAt >= 0) OnReconnecting?.Invoke(_reconnectAttempt);
        }

        void TickBrowser()
        {
            if (_reconnectAt >= 0 && Time.realtimeSinceStartup >= _reconnectAt)
            {
                _reconnectAt = -1;
                if (!_manualClose && AutoReconnect) StartBrowserConnection();
            }
            int current = _browserSocket;
            // Bound work per frame; callbacks may close or replace this connection.
            for (int i = 0; current != 0 && current == _browserSocket && i < 100; i++)
            {
                string item = WebSocketBridge.AIGame_WsPoll(current);
                if (string.IsNullOrEmpty(item)) break;
                if (item[0] == 'O')
                {
                    State = WsState.Connected;
                    _reconnectAttempt = 0;
                    LastError = null;
                    OnConnected?.Invoke();
                }
                else if (item[0] == 'M') OnMessage?.Invoke(item.Substring(1));
                else if (item[0] == 'E' || item[0] == 'C') BrowserDisconnected(item.Substring(1));
            }
            if (_browserSocket != 0 && State != WsState.Connected && Time.realtimeSinceStartup >= _connectDeadline)
                BrowserDisconnected("WebSocket connection timed out");
            for (int i = 0; _browserSocket != 0 && State == WsState.Connected && i < 100 && _outbox.TryDequeue(out var text); i++)
                WebSocketBridge.AIGame_WsSend(_browserSocket, text);
        }
#else
        async Task ConnectInternalAsync()
        {
            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 */ }
                }
            }
        }

        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;
        }
#endif

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