#!/usr/bin/env ruby

# Read a PNG file and create a palette-locked IFF ILBM file
# This allows for multiple files to be exported by other tools
# with the same palette colors, and for the resulting IFF ILBM
# images to contain those pixels, with those colors, with the
# color indexes the same across all images.

# References:
#
# * http://etwright.org/lwsdk/docs/filefmts/ilbm.html

require 'rmagick'

class PixelColor
  attr_reader :r, :g, :b

  def initialize(r:, g:, b:)
    @r = r
    @g = g
    @b = b
  end

  def ==(other)
    r == other.r && g == other.g && b == other.b
  end

  def to_cmap
    [
      r * 16 + r,
      g * 16 + g,
      b * 16 + b
    ]
  end

  def inspect
    "RGB: #{r.to_s(16)},#{g.to_s(16)},#{b.to_s(16)}"
  end
end

DESIRED_PALETTE = [
  [15,0,15],
  [0,0,0],
  [15,15,15],
  [5,0xd,3],
  [0xe, 0xe, 0],
  [0xe, 0xd, 0xa],
  [3, 8, 0xd],
  [4, 2, 2],
  [0xd, 0xe, 0xd],
  [0xe, 0xa, 0xe],
  [8, 1, 2],
  [3, 8, 2],
  [7, 6, 4],
  [5, 5, 5],
  [0xf, 0xf, 8],
  [0xe, 9, 3]  
].map { |p| PixelColor.new(r: p[0], g: p[1], b: p[2]) }

class Bitplane
  def initialize(bits:)
    @bits = bits
  end

  def count
    @bits.count
  end

  def [](idx)
    @bits[idx]
  end
end

# Convert an RMagick image into bitplane data based on
# the provided palette. Raises UnknownPaletteColor if there's a
# color in the image that's not in the palette.
class ImageComponentParser
  # @param image [Magick::Image]
  def self.parse(image:, desired_palette:)
    parser = new(image: image, desired_palette: desired_palette)
    parser.parse!
    parser
  end

  UnknownPaletteColor = Class.new(ArgumentError)

  def initialize(image:, desired_palette:)
    @image = image
    @desired_palette = desired_palette
    @bitplanes = nil
  end

  def bitplane_count
    Math.sqrt(@desired_palette.length).ceil
  end

  def bitplanes
    return @bitplanes if @bitplanes

    bits = []
    # get around duplicate objects
    bitplane_count.times { |i| bits << [] }

    @bytes.each do |byte|
      mask = 1
      bitplane_count.times do |bit|
        bits[bit] << (((byte & mask) == mask) ? 1 : 0)

        mask = mask << 1
      end
    end

    @bitplanes = bits.map { |b| Bitplane.new(bits: b) }
  end

  def parse!
    @bytes = []
    @image.each_pixel do |px|
      pixel_color = PixelColor.new(
        # 4 bit amiga pixel colors
        r: px.red >> 12,
        g: px.green >> 12,
        b: px.blue >> 12
      )

      index = @desired_palette.find_index { |p| p == pixel_color }
      unless index
        raise UnknownPaletteColor, "Palette color not found! Check export: #{pixel_color}"  
      end

      @bytes << index
    end
  end
end

module IFF
  module WithHeader
    def to_bytes
      data_bytes = to_data_bytes
      data_bytes_count = data_bytes.count

      output = self.class.name.split("::").last.unpack("C*") +
        [data_bytes_count].pack("N").unpack("C*") +
        data_bytes

      # pad chunk
      output << 0 if data_bytes_count.odd?

      output
    end
  end

  # Construct an IFF ILBM image from the provided data.
  class ILBMBuilder
    def initialize(
      width:,
      height:,
      palette:,
      bitplanes:,
      row_writer: CompressedRowWriter,
      screen_format: CAMG.PalLowRes
    )
      @width = width
      @height = height
      @palette = palette
      @bitplanes = bitplanes
      @row_writer = row_writer
      @screen_format = screen_format
    end

    def row_writer
      return UncompressedRowWriter

      # here's no reason to use compression if the image is less 
      # than 24 pixels wide
      if @row_writer == CompressedRowWriter && @width < 24
        return UncompressedRowWriter
      end

      @row_writer
    end

    def build
      form = FORM.new

      bmhd = BMHD.new
      bmhd.width = @width
      bmhd.height = @height
      bmhd.planes = Math.sqrt(@palette.length).ceil
      bmhd.compression = row_writer.bmhd_flag

      cmap = CMAP.new(palette: @palette)

      body = BODY.new(
        bitplanes: @bitplanes,
        width: @width,
        row_writer: row_writer
      )

      ilbm = ILBM.new
      ilbm.bmhd = bmhd
      ilbm.cmap = cmap
      ilbm.body = body
      ilbm.camg = @screen_format

      form << ilbm

      form
    end
  end

  # IFF wrapper
  class FORM
    include WithHeader

    def initialize
      @chunks = []
    end

    def <<(chunk)
      @chunks << chunk
    end

    def to_data_bytes
      @chunks.flat_map(&:to_bytes)
    end

    def to_s
      to_bytes.pack("C*")
    end
  end

  # Image type header
  class ILBM
    attr_accessor :bmhd, :cmap, :body, :camg

    def to_bytes
      "ILBM".unpack("C*") + 
      bmhd.to_bytes +
      cmap.to_bytes +
      camg.to_bytes +
      # body is always last
      body.to_bytes
    end
  end

  # Image properties
  class BMHD
    include WithHeader

    ASPECT_RATIO = 0x2C

    attr_accessor :width, :height, :planes, :compression

    def to_data_bytes
      [
        *word_to_bytes(width),
        *word_to_bytes(height),
        0, 0,
        0, 0,
        planes,
        0,
        compression,
        0,
        0, 0,
        ASPECT_RATIO, ASPECT_RATIO,
        *word_to_bytes(width),
        *word_to_bytes(height),
      ]
    end

    private

    def word_to_bytes(word)
      [ word >> 8, word & 0xff ]
    end
  end

  # Color map
  class CMAP
    include WithHeader

    def initialize(palette:)
      @palette = palette
    end

    def to_data_bytes
      @palette.flat_map(&:to_cmap)
    end
  end

  # Amiga-related display properties
  class CAMG
    include WithHeader

    def self.PalLowRes
      new(0x00021000)
    end

    def initialize(setting)
      @setting = setting
    end

    def to_data_bytes
      [@setting].pack("N").unpack("C*")
    end
  end

  module AsRowWriter
    def initialize
      @data = []
    end

    def <<(byte)
      @data << byte
    end

    # Output pixel data is word padded
    def ensure_padded!
      @data << 0 if @data.length.odd?
    end
  end

  # Write out uncompressed bitplane data
  class UncompressedRowWriter
    include AsRowWriter

    def self.bmhd_flag
      0
    end

    def to_bytes
      ensure_padded!

      @data
    end
  end

  # Write out compressed bitplane data using the
  # ByteRun1 run encoding algorithm.
  class CompressedRowWriter
    include AsRowWriter

    def self.bmhd_flag
      1
    end

    def to_bytes
      ensure_padded!

      runs = []

      # discover initial runs
      # format is [byte, run count]
      @data.each do |byte|
        if !runs.last
          runs << [byte, 1]
        elsif runs.last.first == byte
          runs.last[1] += 1
        else
          runs << [byte, 1]
        end
      end

      # categorize and pack literals
      # formats are:
      # * [:repeat, byte_count]
      # * [:literal, byte, ...]
      runs = runs.each_with_object([]) do |(byte, count), out|
        if count > 1
          out << [:repeat, byte, count]
        else
          # pack literals together
          if out.last&.first != :literal
            out << [:literal]
          end

          out.last << byte
        end
      end

      # write out compressed data
      runs.each_with_object([]) do |type, obj|
        case type.first
        when :repeat
          # signed int
          obj << (129 - type[2]) + 128
          obj << type[1]
        when :literal
          # n-1
          obj << type.length - 2
          obj.concat(type[1..])
        end
      end
    end
  end

  # Pixel data
  class BODY
    include WithHeader

    def initialize(bitplanes:, width:, row_writer:)
      @bitplanes = bitplanes
      @width = width
      @row_writer = row_writer
    end

    START_MASK = 1 << 7

    def to_data_bytes
      rows = []

      row_count = @bitplanes[0].count / @width

      row_count.times do |y|
        # interleave the bitplanes
        pos = y * @width

        @bitplanes.each do |bitplane|
          val = 0
          mask = START_MASK

          # the row writer takes care of turning bitplane data into output
          current_row_writer = @row_writer.new

          @width.times do |x|
            val = bitplane[pos + x] * mask | val
            mask >>= 1

            # end of byte
            if mask == 0
              current_row_writer << val

              val = 0
              mask = START_MASK
            end
          end

          unless mask == START_MASK
            current_row_writer << val
          end

          rows.concat(current_row_writer.to_bytes)
        end
      end

      rows
    end
  end
end

image = Magick::Image.read(ARGV[0]).first
parsed = ImageComponentParser.parse(image: image, desired_palette: DESIRED_PALETTE)

form = IFF::ILBMBuilder.new(
  width: image.columns,
  height: image.rows,
  palette: DESIRED_PALETTE,
  bitplanes: parsed.bitplanes,
).build

File.open(ARGV[1], 'wb') { |fh| fh.puts form.to_s }