XFESystem

using System.Buffers.Binary;
using ICSharpCode.SharpZipLib.Checksum;
using ICSharpCode.SharpZipLib.Zip.Compression;

namespace XFESystem.Browser;

/// <summary>Bounded PNG scanlines with managed DEFLATE, including Adam7 and palette transparency.
/// No platform zlib import is required. Color-management metadata and animation are not interpreted.</summary>
internal static class BrowserPngDecoder
{
    internal static bool HasSignature(ReadOnlySpan<byte> bytes) => bytes.Length >= 8 && bytes[..8].SequenceEqual(new byte[] { 137, 80, 78, 71, 13, 10, 26, 10 });
    internal static BrowserCanvasSurface Decode(byte[] bytes, int maximumPixels)
    {
        if (!HasSignature(bytes) || bytes.Length > 16_777_216) throw new IOException("Invalid or oversized PNG input.");
        int width = 0, height = 0, depth = 0, colorType = 0, interlace = 0, chunks = 0, position = 8;
        bool header = false, data = false, dataEnded = false, ended = false;
        byte[] palette = [], transparency = [];
        using var compressed = new MemoryStream();
        while (position < bytes.Length)
        {
            if (++chunks > 4096 || bytes.Length - position < 12) throw new IOException("Invalid PNG chunk framing.");
            uint size = BinaryPrimitives.ReadUInt32BigEndian(bytes.AsSpan(position, 4));
            if (size > bytes.Length - position - 12) throw new IOException("Truncated PNG chunk.");
            int length = (int)size;
            string kind = System.Text.Encoding.ASCII.GetString(bytes, position + 4, 4);
            if (!kind.All(char.IsAsciiLetter)) throw new IOException("Invalid PNG chunk type.");
            var crc = new Crc32(); crc.Update(new ArraySegment<byte>(bytes, position + 4, length + 4));
            if ((uint)crc.Value != BinaryPrimitives.ReadUInt32BigEndian(bytes.AsSpan(position + 8 + length, 4))) throw new IOException("PNG chunk checksum mismatch.");
            var content = bytes.AsSpan(position + 8, length);
            if (!header && kind != "IHDR") throw new IOException("PNG header must be first.");
            if (data && kind != "IDAT") dataEnded = true;
            switch (kind)
            {
                case "IHDR":
                    if (header || length != 13) throw new IOException("Invalid PNG header.");
                    uint w = BinaryPrimitives.ReadUInt32BigEndian(content), h = BinaryPrimitives.ReadUInt32BigEndian(content[4..]);
                    if (w == 0 || h == 0 || w > 32768 || h > 32768 || (long)w * h > maximumPixels)
                        throw new BrowserLimitException("PNG dimensions exceed the decoded pixel budget.");
                    width = (int)w; height = (int)h; depth = content[8]; colorType = content[9]; interlace = content[12];
                    bool validDepth = colorType switch { 0 or 3 => depth is 1 or 2 or 4 or 8 || colorType == 0 && depth == 16,
                        2 or 4 or 6 => depth is 8 or 16, _ => false };
                    if (!validDepth || content[10] != 0 || content[11] != 0 || interlace > 1) throw new NotSupportedException("Unsupported PNG coding method.");
                    header = true; break;
                case "PLTE":
                    if (data || palette.Length != 0 || length == 0 || length > 768 || length % 3 != 0 || colorType is 0 or 4) throw new IOException("Invalid PNG palette.");
                    palette = content.ToArray(); break;
                case "tRNS":
                    if (data || transparency.Length != 0 || colorType is 4 or 6 ||
                        colorType == 0 && length != 2 || colorType == 2 && length != 6 ||
                        colorType == 3 && (palette.Length == 0 || length == 0 || length > palette.Length / 3)) throw new IOException("Invalid PNG transparency.");
                    transparency = content.ToArray(); break;
                case "IDAT":
                    if (dataEnded || colorType == 3 && palette.Length == 0) throw new IOException("Invalid PNG image-data ordering.");
                    compressed.Write(content); data = true; break;
                case "IEND":
                    if (!data || length != 0) throw new IOException("Invalid PNG end marker.");
                    ended = true; break;
                default:
                    if (char.IsAsciiLetterUpper(kind[0])) throw new NotSupportedException("Unsupported critical PNG chunk: " + kind);
                    break;
            }
            position += length + 12;
            if (ended) break;
        }
        if (!ended || position != bytes.Length) throw new IOException("PNG image data is incomplete.");
        if (colorType == 3 && palette.Length / 3 > 1 << depth) throw new IOException("PNG palette exceeds its bit depth.");
        var inflater = new Inflater(); inflater.SetInput(compressed.ToArray());
        var surface = new BrowserCanvasSurface(width, height);
        int channels = colorType switch { 2 => 3, 4 => 2, 6 => 4, _ => 1 };
        int[] startX = interlace == 0 ? [0] : [0, 4, 0, 2, 0, 1, 0];
        int[] startY = interlace == 0 ? [0] : [0, 0, 4, 0, 2, 0, 1];
        int[] stepX = interlace == 0 ? [1] : [8, 8, 4, 4, 2, 2, 1];
        int[] stepY = interlace == 0 ? [1] : [8, 8, 8, 4, 4, 2, 2];
        byte[] filterByte = new byte[1];
        for (int pass = 0; pass < startX.Length; pass++)
        {
            int columns = Math.Max(0, (width - startX[pass] + stepX[pass] - 1) / stepX[pass]);
            int rows = Math.Max(0, (height - startY[pass] + stepY[pass] - 1) / stepY[pass]);
            if (columns == 0 || rows == 0) continue;
            int rowLength = checked((columns * channels * depth + 7) / 8), distance = Math.Max(1, (channels * depth + 7) / 8);
            byte[] row = new byte[rowLength], previous = new byte[rowLength];
            for (int y = 0; y < rows; y++)
            {
                ReadExactly(inflater, filterByte); ReadExactly(inflater, row);
                int filter = filterByte[0]; if (filter > 4) throw new IOException("Invalid PNG row filter.");
                for (int index = 0; index < row.Length; index++)
                {
                    int left = index >= distance ? row[index - distance] : 0, above = previous[index], upperLeft = index >= distance ? previous[index - distance] : 0;
                    int predictor = filter switch { 1 => left, 2 => above, 3 => (left + above) / 2, 4 => Paeth(left, above, upperLeft), _ => 0 };
                    row[index] = unchecked((byte)(row[index] + predictor));
                }
                for (int x = 0; x < columns; x++)
                {
                    int first = Sample(row, x * channels, depth), red, green, blue, alpha = 255;
                    if (colorType == 3)
                    {
                        if (first >= palette.Length / 3) throw new IOException("PNG palette index is outside its palette.");
                        red = palette[first * 3]; green = palette[first * 3 + 1]; blue = palette[first * 3 + 2];
                        if (first < transparency.Length) alpha = transparency[first];
                    }
                    else if (colorType is 0 or 4)
                    {
                        red = green = blue = Expand(first, depth);
                        if (colorType == 4) alpha = Expand(Sample(row, x * channels + 1, depth), depth);
                        else if (transparency.Length == 2 && first == BinaryPrimitives.ReadUInt16BigEndian(transparency)) alpha = 0;
                    }
                    else
                    {
                        int g = Sample(row, x * channels + 1, depth), b = Sample(row, x * channels + 2, depth);
                        red = Expand(first, depth); green = Expand(g, depth); blue = Expand(b, depth);
                        if (colorType == 6) alpha = Expand(Sample(row, x * channels + 3, depth), depth);
                        else if (transparency.Length == 6 && first == BinaryPrimitives.ReadUInt16BigEndian(transparency) &&
                            g == BinaryPrimitives.ReadUInt16BigEndian(transparency.AsSpan(2)) && b == BinaryPrimitives.ReadUInt16BigEndian(transparency.AsSpan(4))) alpha = 0;
                    }
                    surface.PutPixel(startX[pass] + x * stepX[pass], startY[pass] + y * stepY[pass], (uint)(alpha << 24 | red << 16 | green << 8 | blue));
                }
                (row, previous) = (previous, row);
            }
        }
        if (inflater.Inflate(filterByte, 0, 1) != 0) throw new BrowserLimitException("PNG decompression exceeds the declared scanline size.");
        if (!inflater.IsFinished || inflater.RemainingInput != 0) throw new IOException("Invalid PNG compressed stream termination.");
        surface.Commit(); return surface;
    }
    private static void ReadExactly(Inflater inflater, byte[] output)
    {
        for (int offset = 0; offset < output.Length;)
        {
            int count = inflater.Inflate(output, offset, output.Length - offset);
            if (count == 0) throw new IOException("PNG compressed scanlines are truncated or require a dictionary.");
            offset += count;
        }
    }
    private static int Sample(byte[] row, int index, int depth) => depth == 8 ? row[index] : depth == 16
        ? row[index * 2] << 8 | row[index * 2 + 1] : row[index * depth / 8] >> (8 - depth - index * depth % 8) & (1 << depth) - 1;
    private static int Expand(int value, int depth) => depth == 8 ? value : (value * 255 + ((1 << depth) - 1) / 2) / ((1 << depth) - 1);
    private static int Paeth(int left, int above, int upperLeft)
    {
        int target = left + above - upperLeft, a = Math.Abs(target - left), b = Math.Abs(target - above), c = Math.Abs(target - upperLeft);
        return a <= b && a <= c ? left : b <= c ? above : upperLeft;
    }
}