module ChunkyPNG::Canvas::PNGEncoding

def determine_png_encoding(constraints = {})
  encoding = case constraints
    when :fast_rgb;         { :color_mode => ChunkyPNG::COLOR_TRUECOLOR, :compression => Zlib::BEST_SPEED }
    when :fast_rgba;        { :color_mode => ChunkyPNG::COLOR_TRUECOLOR_ALPHA, :compression => Zlib::BEST_SPEED }
    when :best_compression; { :compression => Zlib::BEST_COMPRESSION, :filtering => ChunkyPNG::FILTER_PAETH }
    when :good_compression; { :compression => Zlib::BEST_COMPRESSION, :filtering => ChunkyPNG::FILTER_NONE }
    when :no_compression;   { :compression => Zlib::NO_COMPRESSION }
    when :black_and_white;  { :color_mode => ChunkyPNG::COLOR_GRAYSCALE, :bit_depth => 1 } 
    when Hash; constraints
    else raise ChunkyPNG::Exception, "Unknown encoding preset: #{constraints.inspect}"
  end
  # Do not create a palette when the encoding is given and does not require a palette.
  if encoding[:color_mode]
    if encoding[:color_mode] == ChunkyPNG::COLOR_INDEXED
      self.encoding_palette = self.palette
      encoding[:bit_depth] ||= self.encoding_palette.determine_bit_depth
    else
      encoding[:bit_depth] ||= 8
    end
  else
    self.encoding_palette = self.palette
    suggested_color_mode, suggested_bit_depth = encoding_palette.best_color_settings
    encoding[:color_mode] ||= suggested_color_mode
    encoding[:bit_depth]  ||= suggested_bit_depth
  end
  # Use Zlib's default for compression unless otherwise provided.
  encoding[:compression] ||= Zlib::DEFAULT_COMPRESSION
  encoding[:interlace] = case encoding[:interlace]
    when nil, false, ChunkyPNG::INTERLACING_NONE; ChunkyPNG::INTERLACING_NONE
    when true, ChunkyPNG::INTERLACING_ADAM7;      ChunkyPNG::INTERLACING_ADAM7
    else encoding[:interlace]
  end
  encoding[:filtering] ||= case encoding[:compression]
    when Zlib::BEST_COMPRESSION; ChunkyPNG::FILTER_PAETH
    when Zlib::NO_COMPRESSION..Zlib::BEST_SPEED; ChunkyPNG::FILTER_NONE
    else ChunkyPNG::FILTER_UP
  end
  return encoding
end

def encode_png_image_pass_to_stream(stream, color_mode, bit_depth, filtering)
  start_pos  = stream.bytesize
  pixel_size = Color.pixel_bytesize(color_mode)
  line_width = Color.scanline_bytesize(color_mode, bit_depth, width)
  
  # Determine the filter method
  encode_method = encode_png_pixels_to_scanline_method(color_mode, bit_depth)
  filter_method = case filtering
    when ChunkyPNG::FILTER_SUB;     :encode_png_str_scanline_sub
    when ChunkyPNG::FILTER_UP;      :encode_png_str_scanline_up
    when ChunkyPNG::FILTER_AVERAGE; :encode_png_str_scanline_average
    when ChunkyPNG::FILTER_PAETH;   :encode_png_str_scanline_paeth
    else nil
  end
  
  0.upto(height - 1) do |y|
    stream << send(encode_method, row(y))
  end
  
  # Now, apply filtering if any
  if filter_method
    (height - 1).downto(0) do |y|
      pos = start_pos + y * (line_width + 1)
      prev_pos = (y == 0) ? nil : pos - (line_width + 1)
      send(filter_method, stream, pos, prev_pos, line_width, pixel_size)
    end
  end
end

def encode_png_image_with_interlacing(color_mode, bit_depth = 8, filtering = ChunkyPNG::FILTER_NONE)
  stream = ChunkyPNG::Datastream.empty_bytearray
  0.upto(6) do |pass|
    subcanvas = self.class.adam7_extract_pass(pass, self)
    subcanvas.encoding_palette = encoding_palette
    subcanvas.encode_png_image_pass_to_stream(stream, color_mode, bit_depth, filtering)
  end
  stream
end

def encode_png_image_without_interlacing(color_mode, bit_depth = 8, filtering = ChunkyPNG::FILTER_NONE)
  stream = ChunkyPNG::Datastream.empty_bytearray
  encode_png_image_pass_to_stream(stream, color_mode, bit_depth, filtering)
  stream
end

def encode_png_pixels_to_scanline_grayscale_1bit(pixels)
  chars = []
  pixels.each_slice(8) do |p1, p2, p3, p4, p5, p6, p7, p8|
    chars << ((p1.nil? ? 0 : (p1 & 0x0000ffff) >> 15 << 7) |
              (p2.nil? ? 0 : (p2 & 0x0000ffff) >> 15 << 6) |
              (p3.nil? ? 0 : (p3 & 0x0000ffff) >> 15 << 5) |
              (p4.nil? ? 0 : (p4 & 0x0000ffff) >> 15 << 4) |
              (p5.nil? ? 0 : (p5 & 0x0000ffff) >> 15 << 3) |
              (p6.nil? ? 0 : (p6 & 0x0000ffff) >> 15 << 2) |
              (p7.nil? ? 0 : (p7 & 0x0000ffff) >> 15 << 1) |
              (p8.nil? ? 0 : (p8 & 0x0000ffff) >> 15))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_grayscale_2bit(pixels)
  chars = []
  pixels.each_slice(4) do |p1, p2, p3, p4|
    chars << ((p1.nil? ? 0 : (p1 & 0x0000ffff) >> 14 << 6) |
              (p2.nil? ? 0 : (p2 & 0x0000ffff) >> 14 << 4) |
              (p3.nil? ? 0 : (p3 & 0x0000ffff) >> 14 << 2) |
              (p4.nil? ? 0 : (p4 & 0x0000ffff) >> 14))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_grayscale_4bit(pixels)
  chars = []
  pixels.each_slice(2) do |p1, p2|
    chars << ((p1.nil? ? 0 : ((p1 & 0x0000ffff) >> 12) << 4) | (p2.nil? ? 0 : ((p2 & 0x0000ffff) >> 12)))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_grayscale_8bit(pixels)
  pixels.map { |p| p >> 8 }.pack("xC#{width}")
end

def encode_png_pixels_to_scanline_grayscale_alpha_8bit(pixels)
  pixels.pack("xn#{width}")
end

def encode_png_pixels_to_scanline_indexed_1bit(pixels)
  chars = []
  pixels.each_slice(8) do |p1, p2, p3, p4, p5, p6, p7, p8|
    chars << ((encoding_palette.index(p1) << 7) |
              (encoding_palette.index(p2) << 6) |
              (encoding_palette.index(p3) << 5) |
              (encoding_palette.index(p4) << 4) |
              (encoding_palette.index(p5) << 3) |
              (encoding_palette.index(p6) << 2) |
              (encoding_palette.index(p7) << 1) |
              (encoding_palette.index(p8)))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_indexed_2bit(pixels)
  chars = []
  pixels.each_slice(4) do |p1, p2, p3, p4|
    chars << ((encoding_palette.index(p1) << 6) |
              (encoding_palette.index(p2) << 4) |
              (encoding_palette.index(p3) << 2) |
              (encoding_palette.index(p4)))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_indexed_4bit(pixels)
  chars = []
  pixels.each_slice(2) do |p1, p2|
    chars << ((encoding_palette.index(p1) << 4) | (encoding_palette.index(p2)))
  end
  chars.pack('xC*')
end

def encode_png_pixels_to_scanline_indexed_8bit(pixels)
  pixels.map { |p| encoding_palette.index(p) }.pack("xC#{width}")
end

def encode_png_pixels_to_scanline_method(color_mode, depth)
  encoder_method = case color_mode
    when ChunkyPNG::COLOR_TRUECOLOR;       :"encode_png_pixels_to_scanline_truecolor_#{depth}bit"
    when ChunkyPNG::COLOR_TRUECOLOR_ALPHA; :"encode_png_pixels_to_scanline_truecolor_alpha_#{depth}bit"
    when ChunkyPNG::COLOR_INDEXED;         :"encode_png_pixels_to_scanline_indexed_#{depth}bit"
    when ChunkyPNG::COLOR_GRAYSCALE;       :"encode_png_pixels_to_scanline_grayscale_#{depth}bit"
    when ChunkyPNG::COLOR_GRAYSCALE_ALPHA; :"encode_png_pixels_to_scanline_grayscale_alpha_#{depth}bit"
    else nil
  end
  
  raise ChunkyPNG::NotSupported, "No encoder found for color mode #{color_mode} and #{depth}-bit depth!" unless respond_to?(encoder_method)
  encoder_method
end

def encode_png_pixels_to_scanline_truecolor_8bit(pixels)
  pixels.pack('x' + ('NX' * width))
end

def encode_png_pixels_to_scanline_truecolor_alpha_8bit(pixels)
  pixels.pack("xN#{width}")
end

def encode_png_pixelstream(color_mode = ChunkyPNG::COLOR_TRUECOLOR, bit_depth = 8, interlace = ChunkyPNG::INTERLACING_NONE, filtering = ChunkyPNG::FILTER_NONE)
  if color_mode == ChunkyPNG::COLOR_INDEXED 
    raise ChunkyPNG::ExpectationFailed, "This palette is not suitable for encoding!" if encoding_palette.nil? || !encoding_palette.can_encode?
    raise ChunkyPNG::ExpectationFailed, "This palette has too many colors!" if encoding_palette.size > (1 << bit_depth)
  end
  case interlace
    when ChunkyPNG::INTERLACING_NONE;  encode_png_image_without_interlacing(color_mode, bit_depth, filtering)
    when ChunkyPNG::INTERLACING_ADAM7; encode_png_image_with_interlacing(color_mode, bit_depth, filtering)
    else raise ChunkyPNG::NotSupported, "Unknown interlacing method: #{interlace}!"
  end
end

def encode_png_str_scanline_average(stream, pos, prev_pos, line_width, pixel_size)
  line_width.downto(1) do |i|
    a = (i > pixel_size) ? stream.getbyte(pos + i - pixel_size) : 0
    b = prev_pos ? stream.getbyte(prev_pos + i) : 0
    stream.setbyte(pos + i, (stream.getbyte(pos + i) - ((a + b) >> 1)) & 0xff)
  end
  stream.setbyte(pos, ChunkyPNG::FILTER_AVERAGE)
end

def encode_png_str_scanline_none(stream, pos, prev_pos, line_width, pixel_size)
  # noop - this method shouldn't get called at all.
end

def encode_png_str_scanline_paeth(stream, pos, prev_pos, line_width, pixel_size)
  line_width.downto(1) do |i|
    a = (i > pixel_size) ? stream.getbyte(pos + i - pixel_size) : 0
    b = (prev_pos) ? stream.getbyte(prev_pos + i) : 0
    c = (prev_pos && i > pixel_size) ? stream.getbyte(prev_pos + i - pixel_size) : 0
    p = a + b - c
    pa = (p - a).abs
    pb = (p - b).abs
    pc = (p - c).abs
    pr = (pa <= pb && pa <= pc) ? a : (pb <= pc ? b : c)
    stream.setbyte(pos + i, (stream.getbyte(pos + i) - pr) & 0xff)
  end
  stream.setbyte(pos, ChunkyPNG::FILTER_PAETH)
end      

def encode_png_str_scanline_sub(stream, pos, prev_pos, line_width, pixel_size)
  line_width.downto(1) do |i|
    a = (i > pixel_size) ? stream.getbyte(pos + i - pixel_size) : 0
    stream.setbyte(pos + i, (stream.getbyte(pos + i) - a) & 0xff)
  end
  stream.setbyte(pos, ChunkyPNG::FILTER_SUB)
end

def encode_png_str_scanline_up(stream, pos, prev_pos, line_width, pixel_size)
  line_width.downto(1) do |i|
    b = prev_pos ? stream.getbyte(prev_pos + i) : 0
    stream.setbyte(pos + i, (stream.getbyte(pos + i) - b) & 0xff)
  end
  stream.setbyte(pos, ChunkyPNG::FILTER_UP)
end

def save(filename, constraints = {})
  File.open(filename, 'wb') { |io| write(io, constraints) }
end

def to_blob(constraints = {})
  to_datastream(constraints).to_blob
end

def to_datastream(constraints = {})
  encoding = determine_png_encoding(constraints)
  ds = Datastream.new
  ds.header_chunk = Chunk::Header.new(:width => width, :height => height,
      :color => encoding[:color_mode], :depth => encoding[:bit_depth], :interlace => encoding[:interlace])
  if encoding[:color_mode] == ChunkyPNG::COLOR_INDEXED
    ds.palette_chunk      = encoding_palette.to_plte_chunk
    ds.transparency_chunk = encoding_palette.to_trns_chunk unless encoding_palette.opaque?
  end
  data           = encode_png_pixelstream(encoding[:color_mode], encoding[:bit_depth], encoding[:interlace], encoding[:filtering])
  ds.data_chunks = Chunk::ImageData.split_in_chunks(data, encoding[:compression])
  ds.end_chunk   = Chunk::End.new
  return ds
end

def write(io, constraints = {})
  to_datastream(constraints).write(io)
end