using System.Buffers.Binary;
using System.IO.Compression;
using System.Text;

namespace Uwaa.PNG;

internal static class PngOpener
{
    public static Png Open(Stream stream)
    {
        ArgumentNullException.ThrowIfNull(stream, nameof(stream));

        if (!stream.CanRead)
            throw new ArgumentException($"The provided stream of type {stream.GetType().FullName} was not readable.");

        if (!HasValidHeader(stream))
            throw new ArgumentException($"The provided stream did not start with the PNG header.");

        Span<byte> crc = stackalloc byte[4];
        ImageHeader imageHeader = ReadImageHeader(stream, crc);

        bool hasEncounteredImageEnd = false;

        Palette? palette = null;

        using MemoryStream output = new MemoryStream();
        using MemoryStream memoryStream = new MemoryStream();
        
        while (TryReadChunkHeader(stream, out var header))
        {
            if (hasEncounteredImageEnd)
                break;

            byte[] bytes = new byte[header.Length];
            int read = stream.Read(bytes, 0, bytes.Length);
            if (read != bytes.Length)
                throw new InvalidOperationException($"Did not read {header.Length} bytes for the {header} header, only found: {read}.");

            if (header.IsCritical)
            {
                switch (header.Name)
                {
                    case "PLTE":
                        if (header.Length % 3 != 0)
                            throw new InvalidOperationException($"Palette data must be multiple of 3, got {header.Length}.");

                        // Ignore palette data unless the header.ColorType indicates that the image is paletted.
                        if (imageHeader.ColorType.HasFlag(ColorType.PaletteUsed))
                            palette = new Palette(bytes);

                        break;

                    case "IDAT":
                        memoryStream.Write(bytes, 0, bytes.Length);
                        break;

                    case "IEND":
                        hasEncounteredImageEnd = true;
                        break;

                    default:
                        throw new NotSupportedException($"Encountered critical header {header} which was not recognised.");
                }
            }
            else
            {
                switch (header.Name)
                {
                    case "tRNS":
                        // Add transparency to palette, if the PLTE chunk has been read.
                        palette?.SetAlphaValues(bytes);
                        break;
                }
            }

            read = stream.Read(crc);
            if (read != 4)
                throw new InvalidOperationException($"Did not read 4 bytes for the CRC, only found: {read}.");

            int result = (int)Crc32.Calculate(Encoding.ASCII.GetBytes(header.Name), bytes);
            int crcActual = (crc[0] << 24) + (crc[1] << 16) + (crc[2] << 8) + crc[3];

            if (result != crcActual)
                throw new InvalidOperationException($"CRC calculated {result} did not match file {crcActual} for chunk: {header.Name}.");
        }

        memoryStream.Flush();
        memoryStream.Seek(2, SeekOrigin.Begin);

        using (DeflateStream deflateStream = new DeflateStream(memoryStream, CompressionMode.Decompress)) 
        {
            deflateStream.CopyTo(output);
            deflateStream.Close();
        }

        byte[] bytesOut = output.ToArray();

        (byte bytesPerPixel, byte samplesPerPixel) = Decoder.GetBytesAndSamplesPerPixel(imageHeader);

        bytesOut = Decoder.Decode(bytesOut, imageHeader, bytesPerPixel, samplesPerPixel);

        return new Png(imageHeader, new RawPngData(bytesOut, bytesPerPixel, palette, imageHeader), palette?.HasAlphaValues ?? false);
    }

    static bool HasValidHeader(Stream stream)
    {
        for (int i = 0; i < ImageHeader.ValidationHeader.Length; i++)
            if (stream.ReadByte() != ImageHeader.ValidationHeader[i])
                return false;

        return true;
    }

    static bool TryReadChunkHeader(Stream stream, out ChunkHeader chunkHeader)
    {
        chunkHeader = default;

        long position = stream.Position;
        if (!StreamHelper.TryReadHeaderBytes(stream, out var headerBytes))
            return false;

        int length = BinaryPrimitives.ReadInt32BigEndian(headerBytes);
        string name = Encoding.ASCII.GetString(headerBytes, 4, 4);
        chunkHeader = new ChunkHeader(position, length, name);

        return true;
    }

    static ImageHeader ReadImageHeader(Stream stream, Span<byte> crc)
    {
        if (!TryReadChunkHeader(stream, out var header))
            throw new ArgumentException("The provided stream did not contain a single chunk.");

        if (header.Name != "IHDR")
            throw new ArgumentException($"The first chunk was not the IHDR chunk: {header}.");

        if (header.Length != 13)
            throw new ArgumentException($"The first chunk did not have a length of 13 bytes: {header}.");

        byte[] ihdrBytes = new byte[13];
        int read = stream.Read(ihdrBytes, 0, ihdrBytes.Length);

        if (read != 13)
            throw new InvalidOperationException($"Did not read 13 bytes for the IHDR, only found: {read}.");

        read = stream.Read(crc);
        if (read != 4)
            throw new InvalidOperationException($"Did not read 4 bytes for the CRC, only found: {read}.");

        int width = BinaryPrimitives.ReadInt32BigEndian(ihdrBytes);
        int height = BinaryPrimitives.ReadInt32BigEndian(ihdrBytes[4..]);
        byte bitDepth = ihdrBytes[8];
        byte colorType = ihdrBytes[9];
        byte interlaceMethod = ihdrBytes[12];

        return new ImageHeader(width, height, bitDepth, (ColorType)colorType, (InterlaceMethod)interlaceMethod);
    }
}