#!ruby -Ku


require "zlib"
require "stringio"

class PNG
	class GeneralPNGError < StandardError; end
	class InvalidPNGError < GeneralPNGError; end
	class UnknownCompressionMethodError < GeneralPNGError; end
	
	FILE_SIGNATURE = "\211PNG\r\n\032\n".freeze
	
	attr_reader :data, :width, :height, :bit_depth, :color_type, :compression_method, :filter_method, :interlace_method
	
	def initialize(data)
		data = StringIO.new(data) if data.class == String
		@data = []  
		@chunk = {} #ただのキャッシュ。
		parse(data)
	end
	
	def [](key)
		@chunk[key.to_sym]
	end
	
	def delete(num)
		@data.delete_at(num)
	end
	
	def add_iTXt(keyword, text, language_tag="", translated_keyword="", compression_flag=false, compression_method=0, level=9)
		stext = text.dup
		if compression_flag
			if compression_method == 0
				text = Zlib::Deflate.deflate(text, level)
				if text.size > stext.size
					text = stext
					compression_flag = false
				end
			else
				UnknownCompressionMethodError
			end
		end
		
		compression_flag = compression_flag ? 1 : 0
		
		data = ""
		data <<  keyword
		data << "\000"
		data << [compression_flag].pack("C")
		data << [compression_method].pack("C")
		data << language_tag
		data << "\000"
		data << translated_keyword
		data << "\000"
		data << text
		contents = {:keyword            => keyword,
		            :compression_flag   => compression_flag,
		            :compression_method => compression_method,
		            :language_tag       => language_tag,
		            :translated_keyword => translated_keyword,
		            :text               => stext}
		@data.insert(-2, [:iTXt, data, contents])
		((@chunk[:iTXt] ||= {})[keyword.to_sym] ||= []) << contents
	end
	
	def dump
		ret = ""
		ret << FILE_SIGNATURE
		@data.each do |chunk|
			ret << [chunk[1].length].pack("N*")
			ret << chunk[0].to_s
			ret << chunk[1]
			ret << [Zlib.crc32(chunk[0].to_s + chunk[1])].pack("N*")
		end
		ret
	end
	
	private
	def parse(io)
		io.rewind
#		p io.read(8), FILE_SIGNATURE
		if io.read(8) == FILE_SIGNATURE
			until io.eof?
				length = io.read(4).unpack("N")[0]
				type   = io.read(4)
				data   = io.read(length)
				crc    = io.read(4)#.unpack("H*")
				raise InvalidPNGError if crc != [Zlib.crc32(type + data)].pack("N*")
				case type
					when "IHDR", "PLTE", "tRNS", "tEXt", "zTXt", "iTXt", "tIME", "gAMA", "pHYs", "bKGD", "sBIT"
						__send__("chunk_#{type}", data)
					else
						chunk_default(type, data)
				end
			end
		else
			raise InvalidPNGError
		end
	end
	
	def chunk_IHDR(data)
		@width              = data[0, 4].unpack("N")[0]
		@height             = data[4, 4].unpack("N")[0]
		@bit_depth          = data[8]
		@color_type         = data[9]
		@compression_method = data[10]
		@filter_method      = data[11]
		@interlace_method   = data[12]
		contents            = {:width              => @width,
		                       :height             => @height,
		                       :bit_depth          => @bit_depth,
		                       :color_type         => @color_type,
		                       :compression_method => @compression_method,
		                       :filter_method      => @filter_method,
		                       :interlace_method   => @interlace_method}
		@data << [:IHDR, data, contents]
		
		@chunk[:IHDR] = contents
	end
	
	def chunk_PLTE(data)
		pallet = []
		data.scan(/.../n) do |rgb|
			r = rgb[0]
			g = rgb[1]
			b = rgb[2]
			pallet << {:r => r, :g => g, :b => b}
		end
		
		
		@data << [:PLTE, data, pallet]
		@chunk[:PLTE] = pallet
	end
	
	def chunk_tRNS(data)
		case @color_type
			when 3
				contents = data.unpack("C*")
			when 0
				contents = data[0, 2].unpack("n")[0]
			when 2
				r = data[0, 2].unpack("n")[0]
				g = data[2, 2].unpack("n")[0]
				b = data[4, 2].unpack("n")[0]
				contents = [r, g, b]
		end
		@data << [:tRNS, data, contents]
		@chunk[:tRNS] = contents
	end
	
	def chunk_tEXt(data)
		keyword, text = data.split(/\000/, 2)
		contents = {:keyword => keyword,
		            :text    => text}
		@data << [:tEXt, data, contents]
		((@chunk[:tEXt] ||= {})[keyword.to_sym] ||= []) << contents
	end
	
	def chunk_zTXt(data)
		keyword, text = data.split(/\000/, 2)
		if text[0] == 0
			text = Zlib::Inflate.inflate(text[1..-1])
		else
			raise UnknownCompressionMethodError
		end
		contents = {:keyword => keyword,
		            :text    => text}
		@data << [:zTXt, data, contents]
		((@chunk[:zTXt] ||= {})[keyword.to_sym] ||= []) << contents
	end
	
	def chunk_iTXt(data)
		keyword, rest= data.split(/\000/, 2)
		compression_flag   = rest[0]
		compression_method = rest[1]
		language_tag, translated_keyword, text = rest[2..-1].split(/\000/, 3)
		if compression_flag == 1
			if compression_method == 0
				text = Zlib::Inflate.inflate(text)
			else
				raise UnknownCompressionMethodError
			end
		end
		contents = {:keyword            => keyword,
		            :compression_flag   => compression_flag,
		            :compression_method => compression_method,
		            :language_tag       => language_tag,
		            :translated_keyword => translated_keyword,
		            :text               => text}
		@data << [:iTXt, data, contents]
		((@chunk[:iTXt] ||= {})[keyword.to_sym] ||= []) << contents
	end
	
	def chunk_tIME(data)
		time = []
		time << data[0, 2].unpack("n")[0]
		time << data[2]
		time << data[3]
		time << data[4]
		time << data[5]
		time << data[6]
		time = Time.utc(*time)
		@data << [:tIME, data, time]
		@chunk[:tIME] = time
	end
	
	def chunk_gAMA(data)
		gamma = data.unpack("N")[0]
		gamma /= 100000.0
		@data << [:gAMA, data, gamma]
		@chunk[:gAMA] = gamma
	end
	
	def chunk_bKGD(data)
		case @color_type
			when 3    # pallet color
				bg = data[0]
				contents = bg
			when 0, 4 # gray scale
				bg = data[0, 2].unpack("n")
				contents = bg
			when 2, 6 # true color
				red   = data[0, 2].unpack("n")
				green = data[2, 2].unpack("n")
				blue  = data[4, 2].unpack("n")
				contents = {:red   => red,
				            :green => green,
				            :blue  => blue}
		end
		@data << [:bKGD, data, contents]
		@chunk[:bKGD] = contents
	end
	
	def chunk_pHYs(data)
		x = data[0, 4].unpack("N")[0]
		y = data[4, 4].unpack("N")[0]
		u = data[5]
		contents = {:x    => x,
		            :y    => y,
		            :unit => u}
		@data << [:pHYs, data, contents]
		@chunk[:pHYs] = contents
	end
	
	def chunk_sBIT(data)
		case @color_type
			when 0
				gray  = data[0]
			when 2
				red   = data[0]
				green = data[1]
				blue  = data[2]
			when 3
				red   = data[0]
				green = data[1]
				blue  = data[2]
			when 4
				gray  = data[0]
				alpha = data[1]
			when 6
				red   = data[0]
				green = data[1]
				blue  = data[2]
				alpha = data[3]
		end
		contents = {:gray  => gray || nil,
		            :red   => red || nil,
		            :green => green || nil,
		            :blue  => blue || nil,
		            :alpha => alpha || nil}
		@data << [:sBIT, data, contents]
		@chunk[:sBIT] = contents
	end
	
	def chunk_default(type, data)
		@data << [type.to_sym, data]
		@chunk[type.to_sym] = true
	end
	
	
end


if $0 == __FILE__
	require "pp"

=begin
	png = nil
	File.open("sample.png", "rb") do |f|
		png = PNG.new(f)
		pp png[:zTXt]
		p png.width, png.height
		png.add_iTXt("Comment", "テスト用", "ja-jp", "コメント", true, 0, 9)
		png.add_iTXt("Comment", "1Test String", "en-us", "komento", true, 0, 9)
		png.add_iTXt("Desc", "2テスト用\n日本語 UTF-8", "ja-jp", "コメント", true, 0, 9)
		png.add_iTXt("Comment", "3Test String", "en-us", "komento", true, 0, 9)
	#	p png.data
	end
	
	pp png[:iTXt]
	open("baz.png", "wb") {|f| f.print png.dump}
=end
	
	png = nil
	fn  = "baz.png"
	File.open(fn, "rb") do |f|
		png = PNG.new(f)
	#	pp png[:iTXt]
		
		
		puts "#{fn} #{png.width}x#{png.height} #{png.bit_depth}bit/sample interlace:#{png.interlace_method == 1 ? 'yes' : 'no'}"
		puts "-" * 47
		png.data.each do |chunk|
			attr = []
			sc = chunk[0].to_s
			attr << ( ((sc[0] & 0b00100000) == 0) ? "critical"       : "ancillary"    )
			attr << ( ((sc[1] & 0b00100000) == 0) ? "public"         : "private"      )
			attr << ( ((sc[2] & 0b00100000) == 0) ? ""               : ""             )
			attr << ( ((sc[3] & 0b00100000) == 0) ? "unsafe to copy" : "safe to copy" )
			puts " %4s | %10s  %7s  %14s " % [sc, attr[0], attr[1], attr[3]]
		end
	end
end