# Yoann Guillot, 2005

# Can read/write a PNG file, if it is 8bpp with a palette (optionnally with the 1st entry transparent)
# 
# Unsupported:
# * Adam7 interlaced
# * non-paletted
# * no 8bpp
# * transparency other that only 1st palette fully transparent
#
# Usage:
# require 'libpng'
# img = Png.read('picture.png')
# if img.px(30, 60) == [255, 0, 0]
# 	puts "oh, it's red !"
# end

require 'zlib'

class Png
	# Reads a png stream from string, starting at offset.
	# Raises RuntimeError on error (invalid/unsupported png feature).
	def Png.read(file)
		Png.new.load File.open(file, 'rb') { |fd| fd.read }
	end

	attr_reader :width, :height, :palette, :lines
	attr_accessor :trans
	
	def rd(raw, len)
		@off += len
		raw[(@off-len)...@off]
	end
	
	def decodeint(txt)
		txt.unpack('C*').inject(0) { |int, chr| (int << 8) | chr }
	end
	private :decodeint, :rd
	
	def initialize(width=nil, height=nil)
		@palette = Array.new
		@lines = Array.new
		@trans = nil
		if width and height
			@width = width
			@height = height
			@height.times { @lines << 0.chr * @width }
			@palette[0] = [0, 0, 0]
		end
	end
	
	def load(raw, off = 0)
		@off = off
		sig = rd(raw, 8).unpack('C*')
		if sig != [137, 80, 78, 71, 13, 10, 26, 10]
			raise "Invalid signature: #{raw[off...(off+8)].inspect}"
		end

		compresseddata = ''
		halfbytes = false
		loop do
			len = rd(raw, 4)
			raise "Premature end of file" if not len
			len = decodeint(len)
			chunktype = rd(raw, 4)
			chunk = rd(raw, len)
			crc = decodeint(rd(raw, 4))

			puts "Chunk #{chunktype}: #{len}o" if $DEBUG

			raise "Bad crc %.8X for #{chunktype}" % crc if crc != Zlib.crc32(chunktype+chunk)
			
			# http://www.w3.org/TR/PNG-Chunks.html
			case chunktype
			when 'IHDR'
				@width = decodeint(chunk[0..3])
				@height = decodeint(chunk[4..7])
				bitdepth = decodeint(chunk[8..8])
				@colourtype = decodeint(chunk[9..9])	# bitmask 1: palette, 2: color, 4: alpha
				compression = decodeint(chunk[10..10])
				filter = decodeint(chunk[11..11])
				interlace = decodeint(chunk[12..12])

				puts "#{@width}x#{@height}x#{bitdepth}bpp, colourtype #{@colourtype}, compression #{compression}, filter #{filter}, interlace #{interlace}" if $DEBUG
				
				case @colourtype
				when 2
					@width *= 3
				when 6
					@width *= 4
				end

				if bitdepth == 4
					halfbytes = true
					bitdepth = 8
				end
				raise 'Unhandled PNG format' if bitdepth != 8 or compression != 0 or filter != 0 or interlace != 0
				
			when 'IEND'
				break
			
			when 'PLTE'
				@palette = Array.new
				(len/3).times { |i|
					@palette << [decodeint(chunk[3*i, 1]), decodeint(chunk[3*i+1, 1]), decodeint(chunk[3*i+2, 1])]
				}
				
				puts "loaded #{@palette.length} colors palette" if $DEBUG

			when 'IDAT'
				compresseddata << chunk
			
			when 'tRNS'
				# alpha transparency palette, 0=transparent, defaults to opaque
				if len == 1 and decodeint(chunk[0, 1]) == 0
					@trans = true
				else
					puts "unhandled transparency chunk - ignored #{chunk[0,64].inspect}" if $VERBOSE
				end
			
			else
				puts "ignoring unhandled chunk type #{chunktype} #{chunk[0,64].inspect}" if $VERBOSE
			end
		end
		
		puts 'Inflating image data' if $DEBUG
		img = Zlib::Inflate.inflate(compresseddata)
		lsize = @width
		lsize /= 2 if halfbytes
		puts "#{@width}x#{@height}, #{img.length} bytes#{want = @height*(lsize+1) ; img.length > want ? ", #{img.length-want} missing" : want > img.length ? ", #{want-img.length} too much" : ', correct size'}" if $DEBUG

		png_buftolines(img, halfbytes)

		if @colourtype & 1 == 0 and ENV['PNG_MAKEPALETTE']
			alpha = true if @colourtype & 4 == 4
			@lines.map! { |l| l.gsub(/...#{'.' if alpha}/m) { |px|
				pal = px.unpack('C*')
				@palette |= [pal]
				@palette.index(pal).chr
			}}
		end

		puts 'png load successful' if $DEBUG
		
		self
	end

	def png_buftolines(img, halfbytes)
		@lines = Array.new

		lsize = @width
		lsize /= 2 if halfbytes

		bpp = 3
		@height.times { |i|
			off = i * (lsize+1)
			type = decodeint(img[off, 1])
			line = img[(off+1)..(off+lsize)]
			line = line.unpack('C*') if ?a.kind_of?(String)
			line = line.map { |c| [c >> 4, c & 15] }.flatten if halfbytes
if $DEBUG and $stdout.tty?
  print type
  $stdout.flush
end
			case type
			when 0	# raw
			when 1	# sub
				@width.times { |j|
					sub = (j >= bpp) ? line[j-bpp] : 0
					line[j] = (line[j] + sub) & 255
				}
			when 2	# up
				@width.times { |j|
					lst = @lines.last ? @lines.last[j] : 0
					line[j] = (line[j] + lst) & 255
				}
			when 3	# sub+up/2
				@width.times { |j|
					sub = (j >= bpp) ? line[j-bpp] : 0
					lst = @lines.last ? @lines.last[j] : 0
					line[j] = (line[j] + (sub+lst)/2) & 255
				}
			when 4	# paeth
				@width.times { |j|
					sub = (j >= bpp) ? line[j-bpp] : 0
					lst = @lines.last ? @lines.last[j] : 0
					sl  = @lines.last ? ((j >= bpp) ? @lines.last[j-bpp] : 0) : 0
					p = sub + lst - sl
					pa = (p-sub).abs
					pb = (p-lst).abs
					pc = (p-sl).abs
					paeth = ((pa <= pb and pa <= pc) ? sub : ((pb <= pc) ? lst : sl))
					line[j] = (line[j] + paeth) & 255
				}
			else raise "Unhandled line encoding #{type.inspect}"
			end
			@lines << line
		}
	end

	# Returns the [r, g, b] array for the specified pixel, or nil if it is transparent.
	# Raises RuntimeError on error (invalid x/y/palette entry)
	def px(x, y)
		raise 'bad y' if y >= @height
		l = @lines[y]
		raise 'bad x' if x >= l.length
		i = l[x]
		raise 'bad color index' if i >= @palette.length
		@palette[i] if i > 0 or not @trans
	end

	def encodeint(i)
		((i >> 24) & 255).chr << ((i>>16) & 255) << ((i>>8) & 255) << (i & 255)
	end
	
	def write_chunk(fd, id, chk)
		fd.write encodeint(chk.length)
		fd.write id
		fd.write chk
		fd.write encodeint(Zlib.crc32(id+chk))
	end
	private :encodeint, :write_chunk
	
	# Writes the current image to a new file
	def write(filename, overwrite = false)
		raise 'File exists' if not overwrite and File.exists? filename

		File.open(filename, 'w') { |fd|
		
			fd.write 137.chr << 80 << 78 << 71 << 13 << 10 << 26 << 10
			
			@colourtype ||= 3
			if @palette.empty?
				@width /= (((@colourtype & 4) == 4) ? 4 : 3)
			end
			chk = encodeint(@width) << encodeint(@height) << 8 << @colourtype << 0 << 0 << 0
			write_chunk(fd, 'IHDR', chk)
			
			chk = @palette.map{ |e| e[0].chr << e[1] << e[2] }.join
			write_chunk(fd, 'PLTE', chk) if not @palette.empty?
			
			write_chunk(fd, 'tRNS', 0.chr) if @trans
			
			chk = @lines.map { |l| 0.chr + l }.join
			chk = Zlib::Deflate.deflate(chk)
			write_chunk(fd, 'IDAT', chk)

			write_chunk(fd, 'IEND', '')
		}
	end

	def inspect
		"<Png: #{@width}x#{@height}, #{@palette.length-1} colors>"
	end
end

begin
	require 'metasm'

	Metasm::DynLdr.new_func_c <<EOS
void png_buf(unsigned char *buf, unsigned width, unsigned height)
{
	for (unsigned i=0 ; i<height ; ++i) {
		unsigned off = i*(width+1);

		switch(buf[off]) {
		case 0:
			break;
		case 1:
			for (unsigned j=3 ; j<width ; ++j)
				buf[off+j+1] += buf[off+j-2];
			break;
		case 2:
			if (i > 0)
				for (unsigned j=0 ; j<width ; ++j)
					buf[off+j+1] += buf[off+j-width];
			break;
		case 3:
			for (unsigned j=0 ; j<width ; ++j) {
				unsigned moo = 0;
				if (j >= 3)
					moo += buf[off+j-2];
				if (i > 0)
					moo += buf[off+j-width];
				buf[off+j+1] += moo/2;
			}
			break;
		case 4:
			for (unsigned j=0 ; j<width ; ++j) {
				int a, b, c, p, pa, pb, pc;

				if (j >= 3) a = buf[off+j-2]; else a = 0;
				if (i > 0) b = buf[off+j-width]; else b = 0;
				if (i > 0 && j >= 3) c = buf[off+j-width-3]; else c = 0;

				p = a + b - c;
				pa = p-a; if (pa < 0) pa = -pa;
				pb = p-b; if (pb < 0) pb = -pb;
				pc = p-c; if (pc < 0) pc = -pc;

				if (pa <= pb && pa <= pc)
					buf[off+j+1] += a;
				else if (pb <= pc)
					buf[off+j+1] += b;
				else
					buf[off+j+1] += c;
			}
			break;
		}
		buf[off] = 0;
	}
}
EOS

class Png
	alias png_buftolines_rb png_buftolines
	def png_buftolines(img, halfbytes)
		return png_buftolines_rb(img, halfbytes) if halfbytes

		Metasm::DynLdr.png_buf(img, @width, @height)

		@lines = Array.new
		@height.times { |i|
			off = i * (@width+1)
			line = img[(off+1)..(off+@width)]
			line = line.unpack('C*') if ?a.kind_of?(String)
			@lines << line
		}
	end
end
rescue Exception
end

if __FILE__ == $0
	Png.read(ARGV.shift)
end