#if UGX_BRIDGE || UNITY_EDITOR
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Uigenx.Protocol;
using UnityEngine;

namespace Uigenx.Bridge
{
    public enum UgxBridgeState
    {
        Disconnected = 0,
        Connecting = 1,
        Connected = 2,
        Failed = 3,
    }

    public sealed class UgxBridgeOptions
    {
        public string Url = UgxSettings.DefaultDaemonUrl;
        public string Token = string.Empty;
        public UgxRole Role = UgxRole.Editor;
        public string ProjectId = string.Empty;
        public string UnityVersion = Application.unityVersion;
        public string PackageVersion = UgxBridgeClient.PackageVersion;
        public string Platform = Application.platform.ToString();
        public bool IsPlaying;
        public UgxRenderPipeline RenderPipeline = UgxRenderPipeline.Unknown;
        public List<UgxCapability> Capabilities = new List<UgxCapability>();
    }

    /// <summary>
    /// WebSocket client for the local UI Creator daemon. The socket lives on a background task and
    /// never touches a Unity API; every inbound message is dispatched from <see cref="Pump"/> on
    /// the main thread.
    /// </summary>
    public sealed class UgxBridgeClient : IDisposable
    {
        public const string PackageVersion = "0.1.0";

        /// <summary>Transport keep alive. Protocol v1 has no Unity to daemon ping, so the daemon pings.</summary>
        public const int PingIntervalSeconds = 15;

        /// <summary>Two missed pings close the socket, which triggers a reconnect.</summary>
        public const int SilenceTimeoutSeconds = PingIntervalSeconds * 2;

        public const int MinBackoffMs = 500;
        public const int MaxBackoffMs = 10_000;

        /// <summary>Outbound queue cap. Oldest messages are dropped first and the loss is reported.</summary>
        public const int MaxOutbound = 256;

        const int ReceiveBufferSize = 16 * 1024;

        /// <summary>Hard cap on one inbound message. A bundle larger than this is a protocol error.</summary>
        public const int MaxInboundBytes = 64 * 1024 * 1024;

        readonly UgxBridgeOptions _options;
        readonly ConcurrentQueue<UgxEnvelope> _inbound = new ConcurrentQueue<UgxEnvelope>();
        readonly ConcurrentQueue<string> _outbound = new ConcurrentQueue<string>();
        readonly ConcurrentQueue<LogLine> _logs = new ConcurrentQueue<LogLine>();
        readonly Dictionary<string, Action<UgxEnvelope>> _handlers =
            new Dictionary<string, Action<UgxEnvelope>>(StringComparer.Ordinal);

        readonly object _deltaLock = new object();
        readonly Dictionary<string, Dictionary<string, object>> _pendingDeltas =
            new Dictionary<string, Dictionary<string, object>>(StringComparer.Ordinal);

        readonly SemaphoreSlim _outboundSignal = new SemaphoreSlim(0);

        // Stopwatch rather than a wall clock: liveness only needs a monotonic source, and it is
        // safe to read from the socket task.
        static readonly System.Diagnostics.Stopwatch Clock = System.Diagnostics.Stopwatch.StartNew();

        CancellationTokenSource _cancellation;
        Task _loop;
        int _state = (int)UgxBridgeState.Disconnected;
        int _droppedOutbound;
        int _deltaSeq;
        long _lastInboundMs;
        UgxBridgeState _pumpedState = UgxBridgeState.Disconnected;

        public UgxBridgeClient(UgxBridgeOptions options)
            => _options = options ?? throw new ArgumentNullException(nameof(options));

        public UgxBridgeState State => (UgxBridgeState)Volatile.Read(ref _state);

        public UgxWelcome Welcome { get; private set; }

        public int DroppedOutboundCount => Volatile.Read(ref _droppedOutbound);

        public event Action<UgxBridgeState> StateChanged;

        /// <summary>Called for a message with no registered handler. Unknown types are warned about.</summary>
        public event Action<UgxEnvelope> UnhandledMessage;

        public void Register(string type, Action<UgxEnvelope> handler)
        {
            if (string.IsNullOrEmpty(type)) throw new ArgumentException("type is required", nameof(type));
            _handlers[type] = handler ?? throw new ArgumentNullException(nameof(handler));
        }

        public void Unregister(string type) => _handlers.Remove(type);

        public void Start()
        {
            if (_loop != null && !_loop.IsCompleted) return;
            _cancellation = new CancellationTokenSource();
            _loop = Task.Run(() => RunAsync(_cancellation.Token));
        }

        public void Stop()
        {
            _cancellation?.Cancel();
            _cancellation = null;
            _loop = null;
            SetState(UgxBridgeState.Disconnected);
        }

        public void Dispose()
        {
            Stop();
            _outboundSignal.Dispose();
        }

        public void Send(string type, object body, string id = null, string re = null)
        {
            var envelope = UgxEnvelope.Create(type, body, DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(), id, re);
            Enqueue(UgxJson.Serialize(envelope));
        }

        public void SendError(string code, string message, string re = null)
            => Send(UgxMessageTypes.Error, new UgxErrorBody { Code = code, Message = message }, re: re);

        /// <summary>
        /// Queues one live value. Changes are coalesced by path until the next <see cref="FlushDeltas"/>,
        /// so a value that changes every frame costs one dictionary write, not one message.
        /// </summary>
        public void QueueDelta(string session, string path, object value)
        {
            if (string.IsNullOrEmpty(session) || string.IsNullOrEmpty(path)) return;
            lock (_deltaLock)
            {
                if (!_pendingDeltas.TryGetValue(session, out var map))
                {
                    map = new Dictionary<string, object>(StringComparer.Ordinal);
                    _pendingDeltas[session] = map;
                }
                map[path] = value;
            }
        }

        /// <summary>Sends one `data.delta` per session holding everything queued since the last flush.</summary>
        public void FlushDeltas()
        {
            lock (_deltaLock)
            {
                foreach (var pair in _pendingDeltas)
                {
                    if (pair.Value.Count == 0) continue;

                    var delta = new UgxDataDelta { Session = pair.Key, Seq = ++_deltaSeq };
                    foreach (var change in pair.Value)
                        delta.Changes.Add(new UgxDataChange { Path = change.Key, V = change.Value });
                    pair.Value.Clear();

                    Send(UgxMessageTypes.DataDelta, delta);
                }
            }
        }

        public void ClearDeltas(string session)
        {
            lock (_deltaLock) _pendingDeltas.Remove(session);
        }

        /// <summary>Main thread. Drains logs and inbound messages and dispatches them to handlers.</summary>
        public void Pump()
        {
            while (_logs.TryDequeue(out var line)) line.Write();

            var previous = _pumpedState;
            var current = State;
            if (previous != current)
            {
                _pumpedState = current;
                try
                {
                    StateChanged?.Invoke(current);
                }
                catch (Exception e)
                {
                    Debug.LogError("[ugx] a bridge state listener threw: " + e);
                }
            }

            while (_inbound.TryDequeue(out var envelope)) Dispatch(envelope);
        }

        void Dispatch(UgxEnvelope envelope)
        {
            if (envelope.T == UgxMessageTypes.Welcome)
            {
                Welcome = envelope.BodyAs<UgxWelcome>();
                return;
            }
            if (envelope.T == UgxMessageTypes.Ping)
            {
                Send(UgxMessageTypes.Pong, null, re: envelope.Id);
                return;
            }

            if (_handlers.TryGetValue(envelope.T, out var handler))
            {
                // Handlers are our own code, but a fault must not stop the pump.
                try
                {
                    handler(envelope);
                }
                catch (Exception e)
                {
                    Debug.LogError("[ugx] bridge handler for '" + envelope.T + "' threw: " + e);
                    SendError(UgxErrorCodes.Internal, e.Message, envelope.Id);
                }
                return;
            }

            if (UnhandledMessage != null)
            {
                UnhandledMessage(envelope);
                return;
            }
            Debug.LogWarning("[ugx] ignoring unknown bridge message '" + envelope.T + "'");
        }

        void Enqueue(string json)
        {
            if (_outbound.Count >= MaxOutbound && _outbound.TryDequeue(out _))
                Interlocked.Increment(ref _droppedOutbound);

            _outbound.Enqueue(json);
            try
            {
                _outboundSignal.Release();
            }
            catch (ObjectDisposedException)
            {
                // Disposed while a send was in flight; the message is simply dropped.
            }
        }

        async Task RunAsync(CancellationToken token)
        {
            int backoff = MinBackoffMs;
            var random = new System.Random();

            while (!token.IsCancellationRequested)
            {
                SetState(UgxBridgeState.Connecting);
                ClientWebSocket socket = null;
                try
                {
                    socket = new ClientWebSocket();
                    socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(PingIntervalSeconds);
                    if (!string.IsNullOrEmpty(_options.Token))
                        socket.Options.SetRequestHeader("Authorization", "Bearer " + _options.Token);

                    await socket.ConnectAsync(new Uri(_options.Url), token).ConfigureAwait(false);

                    SetState(UgxBridgeState.Connected);
                    backoff = MinBackoffMs;
                    Volatile.Write(ref _lastInboundMs, Clock.ElapsedMilliseconds);
                    SendHello();

                    using (var pair = CancellationTokenSource.CreateLinkedTokenSource(token))
                    {
                        var send = SendLoopAsync(socket, pair.Token);
                        var receive = ReceiveLoopAsync(socket, pair.Token);
                        await Task.WhenAny(send, receive).ConfigureAwait(false);
                        pair.Cancel();
                    }
                }
                catch (OperationCanceledException)
                {
                    break;
                }
                catch (Exception e)
                {
                    Log(LogType.Log, "[ugx] bridge connection failed: " + e.Message);
                    SetState(UgxBridgeState.Failed);
                }
                finally
                {
                    socket?.Dispose();
                }

                if (token.IsCancellationRequested) break;
                SetState(UgxBridgeState.Disconnected);

                // Exponential backoff with jitter so many editors do not retry in lockstep.
                int jitter = random.Next(0, backoff / 2);
                try
                {
                    await Task.Delay(backoff + jitter, token).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
                backoff = Math.Min(MaxBackoffMs, backoff * 2);
            }

            SetState(UgxBridgeState.Disconnected);
        }

        void SendHello()
        {
            var hello = new UgxHello
            {
                Token = _options.Token ?? string.Empty,
                Role = _options.Role,
                UnityVersion = _options.UnityVersion,
                PackageVersion = _options.PackageVersion,
                ProjectId = _options.ProjectId,
                Platform = _options.Platform,
                IsPlaying = _options.IsPlaying,
                RenderPipeline = _options.RenderPipeline,
                Capabilities = _options.Capabilities,
            };
            Send(UgxMessageTypes.Hello, hello);
        }

        async Task SendLoopAsync(ClientWebSocket socket, CancellationToken token)
        {
            while (!token.IsCancellationRequested && socket.State == WebSocketState.Open)
            {
                await _outboundSignal.WaitAsync(1000, token).ConfigureAwait(false);

                while (_outbound.TryDequeue(out var json))
                {
                    var bytes = Encoding.UTF8.GetBytes(json);
                    await socket.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Text, true, token)
                        .ConfigureAwait(false);
                }
                // One Release happened per queued message but the drain above took them all, so
                // the leftover permits are consumed to keep the loop from spinning.
                while (_outboundSignal.Wait(0)) { }

                long silentMs = Clock.ElapsedMilliseconds - Volatile.Read(ref _lastInboundMs);
                if (silentMs <= SilenceTimeoutSeconds * 1000L) continue;

                Log(LogType.Warning, "[ugx] daemon went silent for " + (silentMs / 1000) + "s, reconnecting");
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "silent", CancellationToken.None)
                    .ConfigureAwait(false);
                return;
            }
        }

        async Task ReceiveLoopAsync(ClientWebSocket socket, CancellationToken token)
        {
            var buffer = new byte[ReceiveBufferSize];
            var message = new MemoryStreamLite();

            while (!token.IsCancellationRequested && socket.State == WebSocketState.Open)
            {
                message.Reset();
                WebSocketReceiveResult result;
                do
                {
                    result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), token).ConfigureAwait(false);
                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, CancellationToken.None)
                            .ConfigureAwait(false);
                        return;
                    }
                    if (message.Length + result.Count > MaxInboundBytes)
                    {
                        Log(LogType.Error, "[ugx] inbound message exceeds " + MaxInboundBytes + " bytes, dropping");
                        return;
                    }
                    message.Write(buffer, result.Count);
                } while (!result.EndOfMessage);

                Volatile.Write(ref _lastInboundMs, Clock.ElapsedMilliseconds);
                Decode(message.ToUtf8String());
            }
        }

        void Decode(string json)
        {
            UgxEnvelope envelope;
            try
            {
                envelope = UgxJson.Deserialize<UgxEnvelope>(json);
            }
            catch (JsonException e)
            {
                Log(LogType.Warning, "[ugx] dropping malformed bridge message: " + e.Message);
                Send(UgxMessageTypes.Error,
                    new UgxErrorBody { Code = UgxErrorCodes.BadRequest, Message = "message is not a protocol envelope" });
                return;
            }

            if (envelope == null || string.IsNullOrEmpty(envelope.T)) return;

            if (envelope.V != UgxJson.ProtocolVersion)
            {
                Send(UgxMessageTypes.Error, new UgxErrorBody
                {
                    Code = UgxErrorCodes.VersionMismatch,
                    Message = "protocol version " + envelope.V + " is not supported, this build speaks " +
                              UgxJson.ProtocolVersion,
                }, re: envelope.Id);
                return;
            }

            _inbound.Enqueue(envelope);
        }

        void SetState(UgxBridgeState state) => Volatile.Write(ref _state, (int)state);

        void Log(LogType type, string message) => _logs.Enqueue(new LogLine(type, message));

        readonly struct LogLine
        {
            readonly LogType _type;
            readonly string _message;

            public LogLine(LogType type, string message)
            {
                _type = type;
                _message = message;
            }

            public void Write()
            {
                switch (_type)
                {
                    case LogType.Error:
                        Debug.LogError(_message);
                        break;
                    case LogType.Warning:
                        Debug.LogWarning(_message);
                        break;
                    default:
                        Debug.Log(_message);
                        break;
                }
            }
        }

        /// <summary>Growable byte buffer reused across messages so receiving does not churn the heap.</summary>
        sealed class MemoryStreamLite
        {
            byte[] _data = new byte[ReceiveBufferSize];
            int _length;

            public int Length => _length;

            public void Reset() => _length = 0;

            public void Write(byte[] source, int count)
            {
                if (_length + count > _data.Length)
                {
                    int capacity = _data.Length;
                    while (capacity < _length + count) capacity *= 2;
                    Array.Resize(ref _data, capacity);
                }
                Buffer.BlockCopy(source, 0, _data, _length, count);
                _length += count;
            }

            public string ToUtf8String() => Encoding.UTF8.GetString(_data, 0, _length);
        }
    }
}
#endif