using System;
using System.IO;
using StbImageWriteSharp;

namespace MultiWheelC.TrajectoryPlanning.Mapping;

/// <summary>编码 RGBA 像素并校验 PNG 文件头、必要块和全部 CRC 的内部写入器。</summary>
internal static class ValidatedPngWriter
{
    private static readonly byte[] Signature = { 137, 80, 78, 71, 13, 10, 26, 10 };
    private static readonly byte[] PhysType = { 112, 72, 89, 115 };

    /// <summary>
    /// 将 RGBA 像素写入经过校验的 PNG 流。
    ///
    /// 参数：rgba 为行主序、每像素四字节的红绿蓝透明度数组；width、height 为像素尺寸；output 为可写目标流。
    /// 返回：无。输入长度不等于 width × height × 4 或 PNG 校验失败时抛出异常。
    /// </summary>
    public static void Write(byte[] rgba, int width, int height, Stream output)
    {
        Write(rgba, width, height, output, 11811u);
    }

    /// <summary>将 RGBA 像素写入带指定物理分辨率元数据的经过校验的 PNG 流。</summary>
    internal static void Write(byte[] rgba, int width, int height, Stream output, uint pixelsPerMeter)
    {
        if (rgba == null || output == null || width <= 0 || height <= 0 || pixelsPerMeter == 0u || rgba.Length != checked(width * height * 4))
            throw new ArgumentException("Invalid RGBA PNG input.");
        byte[] encoded;
        using (var memory = new MemoryStream())
        {
            new ImageWriter().WritePng(rgba, width, height, ColorComponents.RedGreenBlueAlpha, memory);
            encoded = memory.ToArray();
        }
        Validate(encoded);
        const int ihdrEndOffset = 8 + 4 + 4 + 13 + 4;
        using (var outputMemory = new MemoryStream())
        {
            outputMemory.Write(encoded, 0, ihdrEndOffset);
            var phys = new byte[9];
            WriteUInt32BigEndian(phys, 0, pixelsPerMeter);
            WriteUInt32BigEndian(phys, 4, pixelsPerMeter);
            phys[8] = 1;
            WriteChunk(outputMemory, PhysType, phys);
            outputMemory.Write(encoded, ihdrEndOffset, encoded.Length - ihdrEndOffset);
            encoded = outputMemory.ToArray();
        }
        Validate(encoded);
        output.Write(encoded, 0, encoded.Length);
    }

    private static void Validate(byte[] png)
    {
        if (png == null || png.Length < 45) throw new InvalidDataException("PNG is truncated.");
        for (int i = 0; i < Signature.Length; i++) if (png[i] != Signature[i]) throw new InvalidDataException("PNG signature is invalid.");
        int offset = Signature.Length, ihdr = 0, iend = 0;
        while (offset < png.Length)
        {
            if (png.Length - offset < 12) throw new InvalidDataException("PNG chunk header is truncated.");
            uint length = ReadUInt32BigEndian(png, offset);
            long crcOffset = (long)offset + 8L + length;
            if (crcOffset + 4L > png.Length) throw new InvalidDataException("PNG chunk is truncated.");
            int typeOffset = offset + 4;
            uint expected = ReadUInt32BigEndian(png, (int)crcOffset);
            uint actual = ComputeCrc32(png, typeOffset, checked((int)length + 4));
            if (expected != actual) throw new InvalidDataException("PNG chunk CRC is invalid.");
            bool isIhdr = IsType(png, typeOffset, 73, 72, 68, 82);
            bool isIend = IsType(png, typeOffset, 73, 69, 78, 68);
            if (offset == Signature.Length && !isIhdr) throw new InvalidDataException("PNG must start with IHDR.");
            if (isIhdr) ihdr++;
            if (isIend) { iend++; if (crcOffset + 4L != png.Length) throw new InvalidDataException("PNG data follows IEND."); }
            offset = checked((int)crcOffset + 4);
        }
        if (ihdr != 1 || iend != 1) throw new InvalidDataException("PNG must contain exactly one IHDR and IEND.");
    }
    private static bool IsType(byte[] bytes, int offset, byte a, byte b, byte c, byte d) { return bytes[offset] == a && bytes[offset + 1] == b && bytes[offset + 2] == c && bytes[offset + 3] == d; }
    private static uint ReadUInt32BigEndian(byte[] bytes, int offset) { return ((uint)bytes[offset] << 24) | ((uint)bytes[offset + 1] << 16) | ((uint)bytes[offset + 2] << 8) | bytes[offset + 3]; }
    private static void WriteUInt32BigEndian(byte[] bytes, int offset, uint value) { bytes[offset] = (byte)(value >> 24); bytes[offset + 1] = (byte)(value >> 16); bytes[offset + 2] = (byte)(value >> 8); bytes[offset + 3] = (byte)value; }
    private static void WriteChunk(Stream output, byte[] type, byte[] data)
    {
        var length = new byte[4]; WriteUInt32BigEndian(length, 0, (uint)data.Length); output.Write(length, 0, 4); output.Write(type, 0, 4); output.Write(data, 0, data.Length);
        var crcBytes = new byte[4]; WriteUInt32BigEndian(crcBytes, 0, ComputeCrc32(type, 0, type.Length, data)); output.Write(crcBytes, 0, 4);
    }
    private static uint ComputeCrc32(byte[] data, int offset, int count) { return ComputeCrc32(data, offset, count, null); }
    private static uint ComputeCrc32(byte[] first, int offset, int count, byte[] second)
    {
        uint crc = 0xffffffffu;
        for (int i = 0; i < count; i++) crc = UpdateCrc(crc, first[offset + i]);
        if (second != null) for (int i = 0; i < second.Length; i++) crc = UpdateCrc(crc, second[i]);
        return crc ^ 0xffffffffu;
    }
    private static uint UpdateCrc(uint crc, byte value)
    {
        crc ^= value;
        for (int bit = 0; bit < 8; bit++) crc = (crc & 1u) == 0u ? crc >> 1 : 0xedb88320u ^ (crc >> 1);
        return crc;
    }
}