#nullable enable
using System;
using System.IO;
using System.IO.Compression;
using System.Text;

namespace IGP.UnitySDK.MirrorTransport
{
    internal static class IGPMirrorPayloadCompression
    {
        public const byte CurrentVersion = 1;
        public const byte DeflateAlgorithm = 1;
        public const int HeaderBytes = 8;

        public const string ReasonAccepted = "accepted";
        public const string ReasonBelowThreshold = "below_threshold";
        public const string ReasonRatioNotMet = "ratio_not_met";

        public readonly struct Diagnostics
        {
            public Diagnostics(
                string reason,
                int originalBytes,
                int compressedBytes,
                int thresholdBytes,
                float minCompressionRatio,
                float actualCompressionRatio,
                string payloadProfile = "")
            {
                Reason = reason;
                OriginalBytes = originalBytes;
                CompressedBytes = compressedBytes;
                ThresholdBytes = thresholdBytes;
                MinCompressionRatio = minCompressionRatio;
                ActualCompressionRatio = actualCompressionRatio;
                PayloadProfile = payloadProfile;
            }

            public string Reason { get; }
            public int OriginalBytes { get; }
            public int CompressedBytes { get; }
            public int ThresholdBytes { get; }
            public float MinCompressionRatio { get; }
            public float ActualCompressionRatio { get; }
            public string PayloadProfile { get; }
        }

        public static bool TryCompress(
            byte[] payload,
            int thresholdBytes,
            float minCompressionRatio,
            out byte[] compressedPayload)
        {
            return TryCompress(
                payload,
                thresholdBytes,
                minCompressionRatio,
                out compressedPayload,
                out _);
        }

        public static bool TryCompress(
            byte[] payload,
            int thresholdBytes,
            float minCompressionRatio,
            out byte[] compressedPayload,
            out Diagnostics diagnostics,
            bool includePayloadProfile = true)
        {
            compressedPayload = Array.Empty<byte>();
            if (payload == null)
            {
                throw new ArgumentNullException(nameof(payload));
            }

            if (thresholdBytes <= 0 || payload.Length < thresholdBytes)
            {
                diagnostics = new Diagnostics(
                    ReasonBelowThreshold,
                    payload.Length,
                    0,
                    thresholdBytes,
                    minCompressionRatio,
                    1f);
                return false;
            }

            if (minCompressionRatio <= 0f || minCompressionRatio >= 1f)
            {
                throw new ArgumentOutOfRangeException(nameof(minCompressionRatio));
            }

            byte[] compressedBytes = CompressDeflate(payload);
            int candidateBytes = compressedBytes.Length + HeaderBytes;
            float actualCompressionRatio = (float)candidateBytes / payload.Length;
            if (candidateBytes >= payload.Length * minCompressionRatio)
            {
                diagnostics = new Diagnostics(
                    ReasonRatioNotMet,
                    payload.Length,
                    candidateBytes,
                    thresholdBytes,
                    minCompressionRatio,
                    actualCompressionRatio,
                    includePayloadProfile ? BuildPayloadProfile(payload) : string.Empty);
                return false;
            }

            compressedPayload = new byte[candidateBytes];
            compressedPayload[0] = (byte)'I';
            compressedPayload[1] = (byte)'G';
            compressedPayload[2] = CurrentVersion;
            compressedPayload[3] = DeflateAlgorithm;
            WriteInt32BigEndian(compressedPayload, 4, payload.Length);
            Buffer.BlockCopy(compressedBytes, 0, compressedPayload, HeaderBytes, compressedBytes.Length);
            diagnostics = new Diagnostics(
                ReasonAccepted,
                payload.Length,
                compressedPayload.Length,
                thresholdBytes,
                minCompressionRatio,
                actualCompressionRatio);
            return true;
        }

        public static byte[] Decompress(byte[] compressedPayload, int maxOriginalBytes)
        {
            if (compressedPayload == null)
            {
                throw new ArgumentNullException(nameof(compressedPayload));
            }

            if (compressedPayload.Length < HeaderBytes ||
                compressedPayload[0] != (byte)'I' ||
                compressedPayload[1] != (byte)'G')
            {
                throw new InvalidOperationException("Compressed Mirror payload header is invalid.");
            }

            if (compressedPayload[2] != CurrentVersion)
            {
                throw new InvalidOperationException(
                    $"Compressed Mirror payload version {compressedPayload[2]} is not supported.");
            }

            if (compressedPayload[3] != DeflateAlgorithm)
            {
                throw new InvalidOperationException(
                    $"Compressed Mirror payload algorithm {compressedPayload[3]} is not supported.");
            }

            int originalBytes = ReadInt32BigEndian(compressedPayload, 4);
            if (originalBytes <= 0 || originalBytes > maxOriginalBytes)
            {
                throw new InvalidOperationException(
                    $"Compressed Mirror payload original size {originalBytes} is invalid.");
            }

            using var input = new MemoryStream(
                compressedPayload,
                HeaderBytes,
                compressedPayload.Length - HeaderBytes,
                writable: false);
            using var deflate = new DeflateStream(input, CompressionMode.Decompress);
            using var output = new MemoryStream(originalBytes);
            deflate.CopyTo(output);
            byte[] result = output.ToArray();
            if (result.Length != originalBytes)
            {
                throw new InvalidOperationException(
                    $"Compressed Mirror payload decompressed to {result.Length} bytes, expected {originalBytes}.");
            }

            return result;
        }

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

            return output.ToArray();
        }

        private static string BuildPayloadProfile(byte[] payload)
        {
            const int MaxSampleBytes = 64 * 1024;
            int sampleBytes = Math.Min(payload.Length, MaxSampleBytes);
            if (sampleBytes <= 0)
            {
                return string.Empty;
            }

            int[] counts = new int[256];
            int printableBytes = 0;
            int zeroBytes = 0;
            int repeatedAdjacentBytes = 0;
            byte previous = 0;
            for (int i = 0; i < sampleBytes; i++)
            {
                byte value = payload[i];
                counts[value]++;
                if (value == 0)
                {
                    zeroBytes++;
                }

                if (IsPrintableAscii(value))
                {
                    printableBytes++;
                }

                if (i > 0 && value == previous)
                {
                    repeatedAdjacentBytes++;
                }

                previous = value;
            }

            int uniqueBytes = 0;
            double entropy = 0d;
            for (int i = 0; i < counts.Length; i++)
            {
                int count = counts[i];
                if (count == 0)
                {
                    continue;
                }

                uniqueBytes++;
                double probability = count / (double)sampleBytes;
                entropy -= probability * Math.Log(probability, 2d);
            }

            double printableRatio = printableBytes / (double)sampleBytes;
            double zeroRatio = zeroBytes / (double)sampleBytes;
            double repeatedRatio = sampleBytes > 1
                ? repeatedAdjacentBytes / (double)(sampleBytes - 1)
                : 0d;

            return $"sampleBytes={sampleBytes} entropyBitsPerByte={entropy:F2} " +
                   $"printableRatio={printableRatio:F3} zeroRatio={zeroRatio:F3} " +
                   $"repeatedAdjacentRatio={repeatedRatio:F3} uniqueBytes={uniqueBytes} " +
                   $"magic={DetectMagic(payload)} first8Hex={FormatFirstBytes(payload, 8)}";
        }

        private static bool IsPrintableAscii(byte value)
        {
            return value == 9 ||
                   value == 10 ||
                   value == 13 ||
                   (value >= 32 && value <= 126);
        }

        private static string DetectMagic(byte[] payload)
        {
            if (StartsWith(payload, 0x89, 0x50, 0x4E, 0x47))
            {
                return "png";
            }

            if (StartsWith(payload, 0xFF, 0xD8, 0xFF))
            {
                return "jpeg";
            }

            if (StartsWith(payload, 0x50, 0x4B, 0x03, 0x04) ||
                StartsWith(payload, 0x50, 0x4B, 0x05, 0x06) ||
                StartsWith(payload, 0x50, 0x4B, 0x07, 0x08))
            {
                return "zip";
            }

            if (StartsWith(payload, 0x1F, 0x8B))
            {
                return "gzip";
            }

            if (StartsWith(payload, 0x04, 0x22, 0x4D, 0x18))
            {
                return "lz4-frame";
            }

            if (StartsWith(payload, 0x28, 0xB5, 0x2F, 0xFD))
            {
                return "zstd";
            }

            return "unknown";
        }

        private static bool StartsWith(byte[] payload, params byte[] prefix)
        {
            if (payload.Length < prefix.Length)
            {
                return false;
            }

            for (int i = 0; i < prefix.Length; i++)
            {
                if (payload[i] != prefix[i])
                {
                    return false;
                }
            }

            return true;
        }

        private static string FormatFirstBytes(byte[] payload, int count)
        {
            int bytesToFormat = Math.Min(payload.Length, count);
            StringBuilder builder = new StringBuilder(bytesToFormat * 2);
            for (int i = 0; i < bytesToFormat; i++)
            {
                if (i > 0)
                {
                    builder.Append('-');
                }

                builder.Append(payload[i].ToString("X2"));
            }

            return builder.ToString();
        }

        private static void WriteInt32BigEndian(byte[] buffer, int offset, int value)
        {
            if (value <= 0)
            {
                throw new ArgumentOutOfRangeException(nameof(value));
            }

            buffer[offset] = (byte)((value >> 24) & 0xFF);
            buffer[offset + 1] = (byte)((value >> 16) & 0xFF);
            buffer[offset + 2] = (byte)((value >> 8) & 0xFF);
            buffer[offset + 3] = (byte)(value & 0xFF);
        }

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