#nullable enable
using System;
using System.Collections.Generic;

namespace IGP.UnitySDK.Protocol
{
    internal sealed class IGPReliableTransportOptions
    {
        public static IGPReliableTransportOptions Default => new IGPReliableTransportOptions(
            reliableMessageMaxBytes: 512 * 1024,
            reliableChunkMaxBytes: 11_520,
            reliableReassemblyTimeout: TimeSpan.FromSeconds(10),
            reliableMaxInflightMessagesPerConnection: 8,
            reliableMaxReassemblyBufferPerConnection: 4 * 1024 * 1024,
            kcpDataPlanePayloadMaxBytes: 12 * 1024,
            kcpFrameMaxBytes: 16 * 1024);

        public IGPReliableTransportOptions(
            int reliableMessageMaxBytes,
            int reliableChunkMaxBytes,
            TimeSpan reliableReassemblyTimeout,
            int reliableMaxInflightMessagesPerConnection,
            int reliableMaxReassemblyBufferPerConnection,
            int kcpDataPlanePayloadMaxBytes,
            int kcpFrameMaxBytes)
        {
            if (reliableMessageMaxBytes <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(reliableMessageMaxBytes));
            }

            if (reliableChunkMaxBytes <= 0 || reliableChunkMaxBytes > reliableMessageMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(reliableChunkMaxBytes));
            }

            if (reliableReassemblyTimeout <= TimeSpan.Zero)
            {
                throw new ArgumentOutOfRangeException(nameof(reliableReassemblyTimeout));
            }

            if (reliableMaxInflightMessagesPerConnection <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(reliableMaxInflightMessagesPerConnection));
            }

            if (reliableMaxReassemblyBufferPerConnection < reliableChunkMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(reliableMaxReassemblyBufferPerConnection));
            }

            if (kcpDataPlanePayloadMaxBytes <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(kcpDataPlanePayloadMaxBytes));
            }

            if (kcpFrameMaxBytes <= 4 || kcpFrameMaxBytes <= kcpDataPlanePayloadMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(kcpFrameMaxBytes));
            }

            ReliableMessageMaxBytes = reliableMessageMaxBytes;
            ReliableChunkMaxBytes = reliableChunkMaxBytes;
            ReliableReassemblyTimeout = reliableReassemblyTimeout;
            ReliableMaxInflightMessagesPerConnection = reliableMaxInflightMessagesPerConnection;
            ReliableMaxReassemblyBufferPerConnection = reliableMaxReassemblyBufferPerConnection;
            KcpDataPlanePayloadMaxBytes = kcpDataPlanePayloadMaxBytes;
            KcpFrameMaxBytes = kcpFrameMaxBytes;
        }

        public int ReliableMessageMaxBytes { get; }
        public int ReliableChunkMaxBytes { get; }
        public TimeSpan ReliableReassemblyTimeout { get; }
        public int ReliableMaxInflightMessagesPerConnection { get; }
        public int ReliableMaxReassemblyBufferPerConnection { get; }
        public int KcpDataPlanePayloadMaxBytes { get; }
        public int KcpFrameMaxBytes { get; }
    }

    internal static class IGPReliableTransportOptionsNegotiation
    {
        public static IGPReliableTransportOptions Resolve(
            int? reliableMessageMaxBytes,
            int? reliableChunkMaxBytes,
            int? kcpDataPlanePayloadMaxBytes,
            int? kcpFrameMaxBytes,
            IGPReliableTransportOptions? defaults = null)
        {
            var fallback = defaults ?? IGPReliableTransportOptions.Default;

            return new IGPReliableTransportOptions(
                reliableMessageMaxBytes: reliableMessageMaxBytes ?? fallback.ReliableMessageMaxBytes,
                reliableChunkMaxBytes: reliableChunkMaxBytes ?? fallback.ReliableChunkMaxBytes,
                reliableReassemblyTimeout: fallback.ReliableReassemblyTimeout,
                reliableMaxInflightMessagesPerConnection: fallback.ReliableMaxInflightMessagesPerConnection,
                reliableMaxReassemblyBufferPerConnection: fallback.ReliableMaxReassemblyBufferPerConnection,
                kcpDataPlanePayloadMaxBytes: kcpDataPlanePayloadMaxBytes ?? fallback.KcpDataPlanePayloadMaxBytes,
                kcpFrameMaxBytes: kcpFrameMaxBytes ?? fallback.KcpFrameMaxBytes);
        }
    }

    internal sealed class IGPReliableChunk
    {
        public IGPReliableChunk(
            string messageId,
            int chunkIndex,
            int chunkCount,
            int totalBytes,
            uint messageType,
            string targetPlayerId,
            byte[] payload)
        {
            MessageId = messageId ?? throw new ArgumentNullException(nameof(messageId));
            ChunkIndex = chunkIndex;
            ChunkCount = chunkCount;
            TotalBytes = totalBytes;
            MessageType = messageType;
            TargetPlayerId = targetPlayerId ?? string.Empty;
            Payload = payload ?? throw new ArgumentNullException(nameof(payload));
        }

        public string MessageId { get; }
        public int ChunkIndex { get; }
        public int ChunkCount { get; }
        public int TotalBytes { get; }
        public uint MessageType { get; }
        public string TargetPlayerId { get; }
        public byte[] Payload { get; }
    }

    internal sealed class IGPReliableReassembledMessage
    {
        public IGPReliableReassembledMessage(string messageId, uint messageType, string targetPlayerId, byte[] payload)
        {
            MessageId = messageId;
            MessageType = messageType;
            TargetPlayerId = targetPlayerId;
            Payload = payload;
        }

        public string MessageId { get; }
        public uint MessageType { get; }
        public string TargetPlayerId { get; }
        public byte[] Payload { get; }
    }

    internal sealed class IGPReliableExpiredMessage
    {
        public IGPReliableExpiredMessage(string connectionId, string messageId)
        {
            ConnectionId = connectionId;
            MessageId = messageId;
        }

        public string ConnectionId { get; }
        public string MessageId { get; }
    }

    internal static class IGPReliableMessageFragmenter
    {
        public static IReadOnlyList<IGPReliableChunk> Fragment(
            byte[] payload,
            uint messageType,
            string targetPlayerId,
            IGPReliableTransportOptions options,
            string? messageId = null)
        {
            if (payload == null)
            {
                throw new ArgumentNullException(nameof(payload));
            }

            if (options == null)
            {
                throw new ArgumentNullException(nameof(options));
            }

            if (payload.Length == 0)
            {
                throw new ArgumentOutOfRangeException(nameof(payload), "Reliable payload must not be empty.");
            }

            if (payload.Length > options.ReliableMessageMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(payload),
                    $"Reliable payload exceeds ReliableMessageMaxBytes={options.ReliableMessageMaxBytes}.");
            }

            string resolvedMessageId = string.IsNullOrWhiteSpace(messageId)
                ? Guid.NewGuid().ToString("N")
                : messageId;

            int chunkCount = (payload.Length + options.ReliableChunkMaxBytes - 1) / options.ReliableChunkMaxBytes;
            var chunks = new List<IGPReliableChunk>(chunkCount);

            for (int index = 0; index < chunkCount; index++)
            {
                int offset = index * options.ReliableChunkMaxBytes;
                int count = Math.Min(options.ReliableChunkMaxBytes, payload.Length - offset);
                var chunkPayload = new byte[count];
                Buffer.BlockCopy(payload, offset, chunkPayload, 0, count);

                chunks.Add(new IGPReliableChunk(
                    resolvedMessageId,
                    index,
                    chunkCount,
                    payload.Length,
                    messageType,
                    targetPlayerId,
                    chunkPayload));
            }

            return chunks;
        }
    }

    internal sealed class IGPReliableMessageReassembler
    {
        private readonly IGPReliableTransportOptions options;
        private readonly Dictionary<string, ConnectionState> connectionStates = new Dictionary<string, ConnectionState>(StringComparer.Ordinal);

        public IGPReliableMessageReassembler(IGPReliableTransportOptions options)
        {
            this.options = options ?? throw new ArgumentNullException(nameof(options));
        }

        public int BufferedBytes { get; private set; }

        public IGPReliableReassembledMessage? AddChunk(string connectionId, IGPReliableChunk chunk, DateTimeOffset receivedAt)
        {
            if (string.IsNullOrWhiteSpace(connectionId))
            {
                throw new ArgumentException("Connection id is required.", nameof(connectionId));
            }

            ValidateChunk(chunk);

            if (!connectionStates.TryGetValue(connectionId, out var connectionState))
            {
                connectionState = new ConnectionState();
                connectionStates[connectionId] = connectionState;
            }

            if (!connectionState.Messages.TryGetValue(chunk.MessageId, out var messageState))
            {
                if (connectionState.Messages.Count >= options.ReliableMaxInflightMessagesPerConnection)
                {
                    throw new InvalidOperationException(
                        $"Reliable inflight messages exceed ReliableMaxInflightMessagesPerConnection={options.ReliableMaxInflightMessagesPerConnection}.");
                }

                if (connectionState.BufferedBytes + chunk.TotalBytes > options.ReliableMaxReassemblyBufferPerConnection)
                {
                    throw new InvalidOperationException(
                        $"Reliable reassembly buffer exceeds ReliableMaxReassemblyBufferPerConnection={options.ReliableMaxReassemblyBufferPerConnection}.");
                }

                messageState = new MessageState(chunk, receivedAt);
                connectionState.Messages[chunk.MessageId] = messageState;
                connectionState.BufferedBytes += chunk.TotalBytes;
                BufferedBytes += chunk.TotalBytes;
            }

            if (messageState.ReceivedChunks[chunk.ChunkIndex] != null)
            {
                return null;
            }

            messageState.ReceivedChunks[chunk.ChunkIndex] = chunk.Payload;
            messageState.ReceivedCount++;
            messageState.LastUpdated = receivedAt;

            if (messageState.ReceivedCount != messageState.ChunkCount)
            {
                return null;
            }

            var payload = new byte[messageState.TotalBytes];
            int offset = 0;
            for (int index = 0; index < messageState.ChunkCount; index++)
            {
                var part = messageState.ReceivedChunks[index];
                if (part == null)
                {
                    return null;
                }

                Buffer.BlockCopy(part, 0, payload, offset, part.Length);
                offset += part.Length;
            }

            if (offset != payload.Length)
            {
                throw new InvalidOperationException("Reliable chunks do not match declared total length.");
            }

            connectionState.Messages.Remove(chunk.MessageId);
            connectionState.BufferedBytes -= messageState.TotalBytes;
            BufferedBytes -= messageState.TotalBytes;

            if (connectionState.Messages.Count == 0)
            {
                connectionStates.Remove(connectionId);
            }

            return new IGPReliableReassembledMessage(
                chunk.MessageId,
                messageState.MessageType,
                messageState.TargetPlayerId,
                payload);
        }

        public IReadOnlyList<IGPReliableExpiredMessage> RemoveExpired(DateTimeOffset now)
        {
            var expired = new List<IGPReliableExpiredMessage>();
            var emptyConnections = new List<string>();

            foreach (var connectionEntry in connectionStates)
            {
                var expiredMessageIds = new List<string>();
                foreach (var messageEntry in connectionEntry.Value.Messages)
                {
                    if (now - messageEntry.Value.LastUpdated < options.ReliableReassemblyTimeout)
                    {
                        continue;
                    }

                    expired.Add(new IGPReliableExpiredMessage(connectionEntry.Key, messageEntry.Key));
                    expiredMessageIds.Add(messageEntry.Key);
                    connectionEntry.Value.BufferedBytes -= messageEntry.Value.TotalBytes;
                    BufferedBytes -= messageEntry.Value.TotalBytes;
                }

                foreach (var messageId in expiredMessageIds)
                {
                    connectionEntry.Value.Messages.Remove(messageId);
                }

                if (connectionEntry.Value.Messages.Count == 0)
                {
                    emptyConnections.Add(connectionEntry.Key);
                }
            }

            foreach (var connectionId in emptyConnections)
            {
                connectionStates.Remove(connectionId);
            }

            return expired;
        }

        public void Reset()
        {
            connectionStates.Clear();
            BufferedBytes = 0;
        }

        private void ValidateChunk(IGPReliableChunk chunk)
        {
            if (chunk == null)
            {
                throw new ArgumentNullException(nameof(chunk));
            }

            if (string.IsNullOrWhiteSpace(chunk.MessageId))
            {
                throw new ArgumentException("Reliable chunk message id is required.", nameof(chunk));
            }

            if (chunk.ChunkCount <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(chunk), "Reliable chunk count must be positive.");
            }

            if (chunk.ChunkIndex < 0 || chunk.ChunkIndex >= chunk.ChunkCount)
            {
                throw new ArgumentOutOfRangeException(nameof(chunk), "Reliable chunk index is out of range.");
            }

            if (chunk.TotalBytes <= 0 || chunk.TotalBytes > options.ReliableMessageMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(chunk),
                    $"Reliable chunk total bytes exceed ReliableMessageMaxBytes={options.ReliableMessageMaxBytes}.");
            }

            if (chunk.Payload.Length <= 0 || chunk.Payload.Length > options.ReliableChunkMaxBytes)
            {
                throw new ArgumentOutOfRangeException(nameof(chunk),
                    $"Reliable chunk payload exceeds ReliableChunkMaxBytes={options.ReliableChunkMaxBytes}.");
            }
        }

        private sealed class ConnectionState
        {
            public Dictionary<string, MessageState> Messages { get; } = new Dictionary<string, MessageState>(StringComparer.Ordinal);
            public int BufferedBytes { get; set; }
        }

        private sealed class MessageState
        {
            public MessageState(IGPReliableChunk chunk, DateTimeOffset receivedAt)
            {
                ChunkCount = chunk.ChunkCount;
                TotalBytes = chunk.TotalBytes;
                MessageType = chunk.MessageType;
                TargetPlayerId = chunk.TargetPlayerId;
                LastUpdated = receivedAt;
                ReceivedChunks = new byte[chunk.ChunkCount][];
            }

            public int ChunkCount { get; }
            public int TotalBytes { get; }
            public uint MessageType { get; }
            public string TargetPlayerId { get; }
            public DateTimeOffset LastUpdated { get; set; }
            public byte[][] ReceivedChunks { get; }
            public int ReceivedCount { get; set; }
        }
    }
}
