#!/usr/bin/env python
#
# Copyright 2013 Google, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Google Author(s): Behdad Esfahbod
#

import struct
import sys
from io import BytesIO


try:
	basestring  # py2
except NameError:
	basestring = str  # py3


class PNG:

	signature = bytearray ((137,80,78,71,13,10,26,10))

	def __init__ (self, f):

		if isinstance(f, basestring):
			f = open (f, 'rb')

		self.f = f
		self.IHDR = None

	def tell (self):
		return self.f.tell ()

	def seek (self, pos):
		self.f.seek (pos)

	def stream (self):
		return self.f

	def data (self):
		self.seek (0)
		return bytearray (self.f.read ())

	class BadSignature (Exception): pass
	class BadChunk (Exception): pass

	def read_signature (self):
		header = bytearray (self.f.read (8))
		if header != PNG.signature:
			raise PNG.BadSignature
		return PNG.signature

	def read_chunk (self):
		buf = self.f.read (4)
		length = struct.unpack (">I", buf)[0]
		chunk_type = self.f.read (4)
		chunk_data = self.f.read (length)
		if len (chunk_data) != length:
			raise PNG.BadChunk
		crc = self.f.read (4)
		if len (crc) != 4:
			raise PNG.BadChunk
		return (chunk_type, chunk_data, crc)

	def read_IHDR (self):
		(chunk_type, chunk_data, crc) = self.read_chunk ()
		if chunk_type != b"IHDR":
			raise PNG.BadChunk
		#  Width:              4 bytes
		#  Height:             4 bytes
		#  Bit depth:          1 byte
		#  Color type:         1 byte
		#  Compression method: 1 byte
		#  Filter method:      1 byte
		#  Interlace method:   1 byte
		return struct.unpack (">IIBBBBB", chunk_data)

	def read_header (self):
		self.read_signature ()
		self.IHDR = self.read_IHDR ()
		return self.IHDR

	def get_size (self):
		if not self.IHDR:
			pos = self.tell ()
			self.seek (0)
			self.read_header ()
			self.seek (pos)
		return self.IHDR[0:2]

	def filter_chunks (self, chunks):
		self.seek (0);
		out = BytesIO ()
		out.write (self.read_signature ())
		while True:
			chunk_type, chunk_data, crc = self.read_chunk ()
			if chunk_type in chunks:
				out.write (struct.pack (">I", len (chunk_data)))
				out.write (chunk_type)
				out.write (chunk_data)
				out.write (crc)
			if chunk_type == b"IEND":
				break
		return PNG (out)