using System.Buffers;
using System.Runtime.CompilerServices;

namespace AyCode.Core.Compression;

/// <summary>
/// Pure managed LZ4 compressor. Works on all platforms including WASM.
/// Implements LZ4 block format compression.
/// </summary>
public static class Lz4Compressor
{
    private const int HashLog = 12;
    private const int HashTableSize = 1 << HashLog;
    private const int MinMatch = 4;
    private const int MaxInputSize = 0x7E000000; // ~2GB
    private const int MFLimit = 12; // Minimum match finding limit from end
    private const int LastLiterals = 5; // Last literals that cannot be matched
    private const int MatchLengthBits = 4;
    private const int LiteralLengthBits = 4;
    private const int RunMask = (1 << LiteralLengthBits) - 1;
    private const int MatchMask = (1 << MatchLengthBits) - 1;

    /// <summary>
    /// Default chunk size for BlockArray mode (64KB).
    /// </summary>
    public const int DefaultChunkSize = 64 * 1024;

    /// <summary>
    /// Compresses data using LZ4 block format.
    /// </summary>
    /// <param name="source">Source data to compress.</param>
    /// <returns>Compressed data.</returns>
    public static byte[] Compress(ReadOnlySpan<byte> source)
    {
        if (source.Length == 0)
            return [];

        if (source.Length > MaxInputSize)
            throw new ArgumentException($"Input too large. Maximum size is {MaxInputSize} bytes.", nameof(source));

        var maxOutputSize = GetMaxCompressedLength(source.Length);
        var output = new byte[maxOutputSize];
        var compressedLength = CompressCore(source, output);

        if (compressedLength == output.Length)
            return output;

        var result = new byte[compressedLength];
        output.AsSpan(0, compressedLength).CopyTo(result);
        return result;
    }

    /// <summary>
    /// Compresses data using LZ4 block format into the destination buffer.
    /// </summary>
    /// <param name="source">Source data to compress.</param>
    /// <param name="destination">Destination buffer for compressed data.</param>
    /// <returns>Number of bytes written to destination.</returns>
    public static int Compress(ReadOnlySpan<byte> source, Span<byte> destination)
    {
        if (source.Length == 0)
            return 0;

        if (source.Length > MaxInputSize)
            throw new ArgumentException($"Input too large. Maximum size is {MaxInputSize} bytes.", nameof(source));

        var maxOutputSize = GetMaxCompressedLength(source.Length);
        if (destination.Length < maxOutputSize)
            throw new ArgumentException($"Destination buffer too small. Need at least {maxOutputSize} bytes.", nameof(destination));

        return CompressCore(source, destination);
    }

    /// <summary>
    /// Compresses data using LZ4 BlockArray format (chunked compression).
    /// </summary>
    /// <param name="source">Source data to compress.</param>
    /// <param name="chunkSize">Size of each chunk (default 64KB).</param>
    /// <returns>Compressed data with chunk headers.</returns>
    public static byte[] CompressBlockArray(ReadOnlySpan<byte> source, int chunkSize = DefaultChunkSize)
    {
        if (source.Length == 0)
            return [];

        if (chunkSize < 1024)
            throw new ArgumentException("Chunk size must be at least 1024 bytes.", nameof(chunkSize));

        var numChunks = (source.Length + chunkSize - 1) / chunkSize;
        var maxOutputSize = 4 + (numChunks * (8 + GetMaxCompressedLength(chunkSize))); // 4 bytes for total chunks

        using var output = new ArrayPoolWriter(maxOutputSize);

        // Write number of chunks
        output.WriteInt32LittleEndian(numChunks);

        var offset = 0;
        while (offset < source.Length)
        {
            var remaining = source.Length - offset;
            var currentChunkSize = Math.Min(remaining, chunkSize);
            var chunk = source.Slice(offset, currentChunkSize);

            var compressedChunk = Compress(chunk);

            // Write original size
            output.WriteInt32LittleEndian(currentChunkSize);
            // Write compressed size
            output.WriteInt32LittleEndian(compressedChunk.Length);
            // Write compressed data
            output.Write(compressedChunk);

            offset += currentChunkSize;
        }

        return output.ToArray();
    }

    /// <summary>
    /// Gets the maximum compressed length for a given input length.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static int GetMaxCompressedLength(int inputLength)
    {
        // LZ4 worst case: input + (input / 255) + 16
        return inputLength + (inputLength / 255) + 16;
    }

    private static int CompressCore(ReadOnlySpan<byte> source, Span<byte> destination)
    {
        var hashTable = ArrayPool<int>.Shared.Rent(HashTableSize);
        try
        {
            Array.Fill(hashTable, -1, 0, HashTableSize);
            return CompressInternal(source, destination, hashTable);
        }
        finally
        {
            ArrayPool<int>.Shared.Return(hashTable);
        }
    }

    private static int CompressInternal(ReadOnlySpan<byte> source, Span<byte> destination, int[] hashTable)
    {
        var srcLength = source.Length;
        if (srcLength < MFLimit)
        {
            // Too small, just store as literals
            return StoreLiterals(source, destination);
        }

        var srcIndex = 0;
        var dstIndex = 0;
        var anchor = 0;
        var srcEnd = srcLength;
        var mfLimit = srcEnd - MFLimit;
        var matchLimit = srcEnd - LastLiterals;

        // First byte is always literal
        srcIndex++;

        while (srcIndex < mfLimit)
        {
            // Find match
            var hash = GetHash(source, srcIndex);
            var matchIndex = hashTable[hash];
            hashTable[hash] = srcIndex;

            // Check if we have a valid match
            if (matchIndex >= 0 &&
                srcIndex - matchIndex <= 65535 &&
                matchIndex >= anchor &&
                ReadUInt32(source, matchIndex) == ReadUInt32(source, srcIndex))
            {
                // Encode literals before match
                var literalLength = srcIndex - anchor;

                // Calculate offset BEFORE extending the match
                var offset = srcIndex - matchIndex;

                // Find match length (extend the match)
                var matchStart = srcIndex;
                var refIndex = matchIndex;
                
                // Skip the 4 bytes we already matched
                srcIndex += MinMatch;
                refIndex += MinMatch;

                // Extend the match
                while (srcIndex < matchLimit && source[srcIndex] == source[refIndex])
                {
                    srcIndex++;
                    refIndex++;
                }

                var matchLength = srcIndex - matchStart - MinMatch;

                // Encode token
                dstIndex = EncodeSequence(destination, dstIndex, source.Slice(anchor, literalLength), offset, matchLength);

                anchor = srcIndex;
            }
            else
            {
                srcIndex++;
            }
        }

        // Encode remaining literals
        var lastLiterals = srcEnd - anchor;
        if (lastLiterals > 0)
        {
            dstIndex = EncodeLiteralsOnly(destination, dstIndex, source.Slice(anchor, lastLiterals));
        }

        return dstIndex;
    }

    private static int StoreLiterals(ReadOnlySpan<byte> source, Span<byte> destination)
    {
        return EncodeLiteralsOnly(destination, 0, source);
    }

    private static int EncodeSequence(Span<byte> dst, int dstIndex, ReadOnlySpan<byte> literals, int offset, int matchLength)
    {
        var literalLength = literals.Length;

        // Encode token
        var token = dstIndex++;

        // Encode literal length
        if (literalLength >= RunMask)
        {
            dst[token] = (byte)(RunMask << MatchLengthBits);
            var remaining = literalLength - RunMask;
            while (remaining >= 255)
            {
                dst[dstIndex++] = 255;
                remaining -= 255;
            }
            dst[dstIndex++] = (byte)remaining;
        }
        else
        {
            dst[token] = (byte)(literalLength << MatchLengthBits);
        }

        // Copy literals
        literals.CopyTo(dst.Slice(dstIndex));
        dstIndex += literalLength;

        // Encode offset (little-endian)
        dst[dstIndex++] = (byte)offset;
        dst[dstIndex++] = (byte)(offset >> 8);

        // Encode match length
        if (matchLength >= MatchMask)
        {
            dst[token] |= MatchMask;
            var remaining = matchLength - MatchMask;
            while (remaining >= 255)
            {
                dst[dstIndex++] = 255;
                remaining -= 255;
            }
            dst[dstIndex++] = (byte)remaining;
        }
        else
        {
            dst[token] |= (byte)matchLength;
        }

        return dstIndex;
    }

    private static int EncodeLiteralsOnly(Span<byte> dst, int dstIndex, ReadOnlySpan<byte> literals)
    {
        var literalLength = literals.Length;

        // Encode token (no match)
        var token = dstIndex++;

        // Encode literal length
        if (literalLength >= RunMask)
        {
            dst[token] = (byte)(RunMask << MatchLengthBits);
            var remaining = literalLength - RunMask;
            while (remaining >= 255)
            {
                dst[dstIndex++] = 255;
                remaining -= 255;
            }
            dst[dstIndex++] = (byte)remaining;
        }
        else
        {
            dst[token] = (byte)(literalLength << MatchLengthBits);
        }

        // Copy literals
        literals.CopyTo(dst.Slice(dstIndex));
        dstIndex += literalLength;

        return dstIndex;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static int GetHash(ReadOnlySpan<byte> data, int index)
    {
        var value = ReadUInt32(data, index);
        return (int)((value * 2654435761u) >> (32 - HashLog));
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static uint ReadUInt32(ReadOnlySpan<byte> data, int index)
    {
        return (uint)(data[index] | (data[index + 1] << 8) | (data[index + 2] << 16) | (data[index + 3] << 24));
    }

    /// <summary>
    /// Helper class for writing to a pooled array.
    /// </summary>
    private sealed class ArrayPoolWriter : IDisposable
    {
        private byte[] _buffer;
        private int _position;

        public ArrayPoolWriter(int initialCapacity)
        {
            _buffer = ArrayPool<byte>.Shared.Rent(initialCapacity);
            _position = 0;
        }

        public void WriteInt32LittleEndian(int value)
        {
            EnsureCapacity(4);
            _buffer[_position++] = (byte)value;
            _buffer[_position++] = (byte)(value >> 8);
            _buffer[_position++] = (byte)(value >> 16);
            _buffer[_position++] = (byte)(value >> 24);
        }

        public void Write(ReadOnlySpan<byte> data)
        {
            EnsureCapacity(data.Length);
            data.CopyTo(_buffer.AsSpan(_position));
            _position += data.Length;
        }

        private void EnsureCapacity(int additionalBytes)
        {
            if (_position + additionalBytes > _buffer.Length)
            {
                var newBuffer = ArrayPool<byte>.Shared.Rent(_buffer.Length * 2);
                _buffer.AsSpan(0, _position).CopyTo(newBuffer);
                ArrayPool<byte>.Shared.Return(_buffer);
                _buffer = newBuffer;
            }
        }

        public byte[] ToArray()
        {
            var result = new byte[_position];
            _buffer.AsSpan(0, _position).CopyTo(result);
            return result;
        }

        public void Dispose()
        {
            ArrayPool<byte>.Shared.Return(_buffer);
        }
    }
}
