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

namespace IGP.Multiplayer.Mirror
{
    internal static class IGPMirrorApplicationFrame
    {
        public const byte CurrentVersion = 1;
        public const byte CompressedFlag = 1;
        public const int HeaderBytes = 12;
        public const int PayloadLengthBytes = 4;
        public const int MaxPayloadCount = 1024;
        public const int MaxRawBodyBytes = 4 * 1024 * 1024;

        public readonly struct EncodedFrame
        {
            public EncodedFrame(byte[] wireData, int payloadCount, int rawBytes, bool compressed)
            {
                WireData = wireData;
                PayloadCount = payloadCount;
                RawBytes = rawBytes;
                Compressed = compressed;
            }

            public byte[] WireData { get; }
            public int PayloadCount { get; }
            public int RawBytes { get; }
            public bool Compressed { get; }
        }

        public static EncodedFrame Encode(IReadOnlyList<byte[]> payloads)
        {
            if (payloads == null)
            {
                throw new ArgumentNullException(nameof(payloads));
            }

            if (payloads.Count == 0 || payloads.Count > MaxPayloadCount)
            {
                throw new ArgumentOutOfRangeException(nameof(payloads));
            }

            int rawBodyBytes = 0;
            for (int i = 0; i < payloads.Count; i++)
            {
                byte[] payload = payloads[i] ?? throw new ArgumentException("Payload cannot be null.", nameof(payloads));
                rawBodyBytes = checked(rawBodyBytes + PayloadLengthBytes + payload.Length);
            }

            if (rawBodyBytes > MaxRawBodyBytes)
            {
                throw new InvalidOperationException($"Application frame body is too large: {rawBodyBytes} bytes.");
            }

            var rawBody = new byte[rawBodyBytes];
            int offset = 0;
            for (int i = 0; i < payloads.Count; i++)
            {
                byte[] payload = payloads[i];
                WriteUInt32LittleEndian(rawBody, offset, (uint)payload.Length);
                offset += PayloadLengthBytes;
                Buffer.BlockCopy(payload, 0, rawBody, offset, payload.Length);
                offset += payload.Length;
            }

            byte[] compressedBody = Compress(rawBody);
            bool useCompression = compressedBody.Length < rawBody.Length;
            byte[] body = useCompression ? compressedBody : rawBody;
            var wire = new byte[HeaderBytes + body.Length];
            wire[0] = (byte)'I';
            wire[1] = (byte)'G';
            wire[2] = (byte)'M';
            wire[3] = (byte)'F';
            wire[4] = CurrentVersion;
            wire[5] = useCompression ? CompressedFlag : (byte)0;
            WriteUInt16LittleEndian(wire, 6, (ushort)payloads.Count);
            WriteUInt32LittleEndian(wire, 8, (uint)rawBody.Length);
            Buffer.BlockCopy(body, 0, wire, HeaderBytes, body.Length);
            return new EncodedFrame(wire, payloads.Count, rawBody.Length, useCompression);
        }

        public static IReadOnlyList<byte[]> Decode(byte[] wireData)
        {
            if (wireData == null)
            {
                throw new ArgumentNullException(nameof(wireData));
            }

            if (wireData.Length < HeaderBytes ||
                wireData[0] != (byte)'I' || wireData[1] != (byte)'G' ||
                wireData[2] != (byte)'M' || wireData[3] != (byte)'F')
            {
                throw new InvalidOperationException("Mirror application frame header is invalid.");
            }

            if (wireData[4] != CurrentVersion)
            {
                throw new InvalidOperationException($"Mirror application frame version {wireData[4]} is not supported.");
            }

            byte flags = wireData[5];
            if ((flags & ~CompressedFlag) != 0)
            {
                throw new InvalidOperationException($"Mirror application frame flags {flags} are invalid.");
            }

            int payloadCount = ReadUInt16LittleEndian(wireData, 6);
            int rawBodyBytes = checked((int)ReadUInt32LittleEndian(wireData, 8));
            if (payloadCount <= 0 || payloadCount > MaxPayloadCount ||
                rawBodyBytes < payloadCount * PayloadLengthBytes || rawBodyBytes > MaxRawBodyBytes)
            {
                throw new InvalidOperationException("Mirror application frame sizes are invalid.");
            }

            byte[] body = (flags & CompressedFlag) != 0
                ? Decompress(wireData, HeaderBytes, wireData.Length - HeaderBytes, rawBodyBytes)
                : CopyBody(wireData, rawBodyBytes);

            var payloads = new List<byte[]>(payloadCount);
            int offset = 0;
            for (int i = 0; i < payloadCount; i++)
            {
                if (offset + PayloadLengthBytes > body.Length)
                {
                    throw new InvalidOperationException("Mirror application frame payload length is truncated.");
                }

                int payloadBytes = checked((int)ReadUInt32LittleEndian(body, offset));
                offset += PayloadLengthBytes;
                if (payloadBytes < 0 || offset + payloadBytes > body.Length)
                {
                    throw new InvalidOperationException("Mirror application frame payload is truncated.");
                }

                var payload = new byte[payloadBytes];
                Buffer.BlockCopy(body, offset, payload, 0, payloadBytes);
                offset += payloadBytes;
                payloads.Add(payload);
            }

            if (offset != body.Length)
            {
                throw new InvalidOperationException("Mirror application frame contains trailing bytes.");
            }

            return payloads;
        }

        private static byte[] CopyBody(byte[] wireData, int expectedBytes)
        {
            if (wireData.Length - HeaderBytes != expectedBytes)
            {
                throw new InvalidOperationException("Mirror application frame body length is invalid.");
            }

            var body = new byte[expectedBytes];
            Buffer.BlockCopy(wireData, HeaderBytes, body, 0, expectedBytes);
            return body;
        }

        private static byte[] Compress(byte[] body)
        {
            using var output = new MemoryStream();
            using (var deflate = new DeflateStream(output, CompressionLevel.Fastest, leaveOpen: true))
            {
                deflate.Write(body, 0, body.Length);
            }

            return output.ToArray();
        }

        private static byte[] Decompress(byte[] wireData, int offset, int count, int expectedBytes)
        {
            using var input = new MemoryStream(wireData, offset, count, writable: false);
            using var deflate = new DeflateStream(input, CompressionMode.Decompress);
            using var output = new MemoryStream(expectedBytes);
            var buffer = new byte[8192];
            while (true)
            {
                int read = deflate.Read(buffer, 0, Math.Min(buffer.Length, expectedBytes - (int)output.Length + 1));
                if (read == 0)
                {
                    break;
                }

                output.Write(buffer, 0, read);
                if (output.Length > expectedBytes)
                {
                    throw new InvalidOperationException("Mirror application frame expands beyond its declared size.");
                }
            }

            byte[] body = output.ToArray();
            if (body.Length != expectedBytes)
            {
                throw new InvalidOperationException("Mirror application frame decompressed length is invalid.");
            }

            return body;
        }

        private static void WriteUInt16LittleEndian(byte[] buffer, int offset, ushort value)
        {
            buffer[offset] = (byte)(value & 0xff);
            buffer[offset + 1] = (byte)(value >> 8);
        }

        private static int ReadUInt16LittleEndian(byte[] buffer, int offset)
        {
            return buffer[offset] | (buffer[offset + 1] << 8);
        }

        private static void WriteUInt32LittleEndian(byte[] buffer, int offset, uint value)
        {
            buffer[offset] = (byte)(value & 0xff);
            buffer[offset + 1] = (byte)((value >> 8) & 0xff);
            buffer[offset + 2] = (byte)((value >> 16) & 0xff);
            buffer[offset + 3] = (byte)((value >> 24) & 0xff);
        }

        private static uint ReadUInt32LittleEndian(byte[] buffer, int offset)
        {
            return buffer[offset] |
                   ((uint)buffer[offset + 1] << 8) |
                   ((uint)buffer[offset + 2] << 16) |
                   ((uint)buffer[offset + 3] << 24);
        }
    }
}
