ApngWriter.kt

import java.awt.Rectangle
import java.io.Closeable
import java.io.IOException
import java.util.zip.CRC32
import okio.Buffer
import okio.BufferedSink
import okio.BufferedSource
import okio.ByteString
import okio.ByteString.Companion.encodeUtf8
import okio.ByteString.Companion.toByteString
import okio.FileHandle
import okio.FileSystem
import okio.Path
import okio.Path.Companion.toPath
import okio.buffer

/**
 * Combine a list of PNGs into an APNG, without decoding any image data.
 *
 * Refs: https://www.w3.org/TR/png/#4Concepts.Encoding, https://wiki.mozilla.org/APNG_Specification
 */
internal class ApngWriter(
  private val fileSystem: FileSystem = FileSystem.SYSTEM,
) {
  private val crcEngine = CRC32()

  fun writePng(
    path: Path,
    fps: Int,
    maxWidth: Int,
    maxHeight: Int,
    inputPaths: List<Path>,
  ) {
    val rectangle = Rectangle(maxWidth, maxHeight)

    fileSystem.write(path) {
      writePngSignature()

      var sequenceNumber = 0
      for (i in inputPaths.indices) {
        val fileHandle = fileSystem.openReadOnly(inputPaths[i])
        val imageChunks = PngReader(fileHandle).use {
          it.readChunks()
        }

        for (chunk in imageChunks) {
          when (chunk.chunkId) {
            IDAT.bytes -> {
              if (i == 0) {
                // Prepend ACTL And FCTL chunks before the first IDAT chunk.
                writeACTL(inputPaths.size, 0)
                writeFCTL(sequenceNumber++, rectangle, fps)
                writeChunk(chunk)
              } else {
                // Prepend an FCTL chunk before each subsequent image's IDAT chunk.
                writeFCTL(sequenceNumber++, rectangle, fps)

                // Convert the IDAT chunk into an FDAT chunk.
                writeChunk(Header(FDAT.bytes)) {
                  writeInt(sequenceNumber++)
                  writeAll(chunk.data)
                }
              }
            }

            // Forward only the first IHDR.
            IHDR.bytes -> {
              if (i == 0) {
                writeChunk(chunk)
              }
            }

            // Don't expect any APNG chunks in input.
            FCTL.bytes, FDAT.bytes, ACTL.bytes -> error("unexpected")

            // Skip the IEND chunk.
            IEND.bytes -> Unit

            // Forward all other chunks as-is, including IHDR.
            else -> writeChunk(chunk)
          }
        }
      }

      writeIEND()
    }
  }

  private fun BufferedSink.writePngSignature() {
    write(PNG_SIG)
  }

  private fun BufferedSink.writeACTL(frameCount: Int, loopCount: Int) {
    writeChunk(ACTL) {
      writeInt(frameCount)
      writeInt(loopCount)
    }
  }

  private fun BufferedSink.writeFCTL(
    sequenceNumber: Int,
    rectangle: Rectangle,
    fps: Int,
  ) {
    writeChunk(FCTL) {
      writeInt(sequenceNumber)
      writeInt(rectangle.width)
      writeInt(rectangle.height)
      writeInt(rectangle.x)
      writeInt(rectangle.y)
      writeShort(1) // Delay Numerator
      writeShort(fps) // Delay Denominator
      writeByte(0) // Dispose Operation
      writeByte(0) // Blend Operation
    }
  }

  private fun BufferedSink.writeIEND() {
    writeChunk(IEND) { }
  }

  private fun BufferedSink.writeChunk(header: Header, data: Buffer.() -> Unit) {
    val buffer = Buffer().apply(data)
    crcEngine.reset()
    crcEngine.update(header.bytes.data, 0, 4)
    val dataBytes = buffer.readByteArray()
    if (dataBytes.isNotEmpty()) crcEngine.update(dataBytes, 0, dataBytes.size)
    val crc = crcEngine.value.toInt()

    writeInt(dataBytes.size)
    write(header.bytes)
    write(dataBytes)
    writeInt(crc)
  }

  private fun BufferedSink.writeChunk(chunk: Chunk) {
    writeChunk(Header(chunk.chunkId)) {
      writeAll(chunk.data)
    }
  }
}

data class Header(val bytes: ByteString)

val PNG_SIG = byteArrayOf(-119, 80, 78, 71, 13, 10, 26, 10).toByteString()
val IHDR = Header("IHDR".encodeUtf8())
val ACTL = Header("acTL".encodeUtf8())
val FCTL = Header("fcTL".encodeUtf8())
val IDAT = Header("IDAT".encodeUtf8())
val FDAT = Header("fdAT".encodeUtf8())
val IEND = Header("IEND".encodeUtf8())


internal class PngReader(
  private val fileHandle: FileHandle,
) : Closeable {
  private val crcEngine = CRC32()
  private val source = fileHandle.source().buffer()

  fun readChunks(): List<Chunk> {
    val result = mutableListOf<Chunk>()

    if (!source.rangeEquals(0L, PNG_SIG)) {
      throw IOException(
        "Missing valid PNG header " +
          "expected: [${PNG_SIG.hex()}] " +
          "actual: [${source.readByteString(PNG_SIG.size.toLong()).hex()}]",
      )
    }

    source.skip(PNG_SIG.size.toLong())

    while (!source.exhausted()) {
      result += source.readChunk()
    }

    return result
  }

  override fun close() {
    fileHandle.close()
    source.close()
  }

  private fun BufferedSource.readChunk(): Chunk {
    val dataLength = readInt()
    val chunkId = readByteString(4L)
    val data = readByteArray(dataLength.toLong())
    val crc = readInt()

    val chunk = Chunk(
      dataLength = dataLength,
      chunkId = chunkId,
      data = Buffer().apply { write(data) },
      crc = crc,
    )

    crcEngine.reset()
    crcEngine.update(chunkId.data, 0, 4)
    if (dataLength > 0) crcEngine.update(data, 0, dataLength)

    if (crcEngine.value.toInt() != crc) {
      throw IOException("CRC Mismatch decoding ${chunkId.utf8()}, invalid data")
    }

    return chunk
  }
}

data class Chunk(
  val dataLength: Int,
  val chunkId: ByteString,
  val data: Buffer,
  val crc: Int,
)

fun main() {
  ApngWriter().writePng(
    "/Users/jwilson/Desktop/hdr/parrot-kotlin-2.png".toPath(),
    30,
    128,
    128,
    listOf(
      "/Users/jwilson/Desktop/hdr/frame1.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame2.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame3.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame4.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame5.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame6.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame7.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame8.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame9.png".toPath(),
      "/Users/jwilson/Desktop/hdr/frame10.png".toPath(),
    ),
  )
}