"""
module pngfile

A more Python-like view of a png file than my previous png.py.  At the moment,
it can do the basic read operations, but cannot write.

For the most part, based on the spec, available at http://www.w3.org/TR/PNG/

Author: Tim Hatch
"""

class PngError(Exception):
    pass

class PngEofError(PngError):
    pass
class PngMagicError(PngError):
    pass
class PngSeekError(PngError):
    pass

class PngChunkError(PngError):
    pass
class PngChunkCrcError(PngChunkError):
    pass
class PngChunkIncompleteError(PngChunkError):
    pass

import struct, zlib

class PngChunk:
    size = 0
    type = "    "
    data = ""
    crc = "    "
    def __init__(self, size=0, type="    ", data="", crc=None):
        self.size = size
        self.type = type
        self.data = data
        if crc is None:
            self.crc = self.getcrc()
            self.size = len(data)
        else:
            self.crc = crc
    
    def getcrc(self):
        if len(self.type) != 4:
            raise PngChunkError, "Type must be four bytes"
        return struct.pack("!L", zlib.crc32(self.type + self.data))
    
    def crcgood(self):
        return self.getcrc() == self.crc
    def __repr__(self):
        return "PngChunk: %s of size %5d, crc %02x %02x %02x %02x" % ((self.type, self.size,)+tuple(map(ord, tuple(self.crc))))
    
class PngFile:
    magic = "\x89\x50\x4e\x47\x0d\x0a\x1a\x0a"
    position = 0
    file = None

    def tell(self):
        """
        Returns the current byte position in the file (i.e. where the next
        read will occur)
        """
        return self.position
    
    def close(self):
        self.file.close()

class PngFileWriter(PngFile):
    def __init__(self, f):
        try:
            f.seek(0)
        except:
            raise PngSeekError, "Can't rewind file"
        self.file = f
        self._write(PngFile.magic)
    
    def _write(self, data):
        self.file.write(data)
    
    def write(self, c):
        if c.size != len(c.data):
            raise PngChunkError, "Chunk size doesn't match chunk data size"
        self._write(struct.pack("!L", c.size))
        if len(c.type) != 4:
            raise PngChunkError, "Type must be 4 bytes"
        self._write(c.type)
        self._write(c.data)
        self._write(c.getcrc())
    
class PngFileReader(PngFile):
    def __init__(self, f, testcrc=False):
        try:
            f.seek(0)
        except:
            raise PngSeekError, "Can't rewind file"
        self.file = f
        self.testcrc = testcrc
        header = self._read(8)
        
        if header != self.magic:
            raise PngMagicError, "Unexpected header, %s" % header
        # At this point, we've got the file pointer ready for reading the first chunk
    
    def _read(self, num, failoneof=True):
        """
        Helper function to read bytes, increment internal position, and error
        if too few bytes are read.  Thank goodness PNG is such a predictable
        stream.
        
        Internal-only, please don't call this without subclassing.
        """
        s = self.file.read(num)
        if len(s) < num:
            if failoneof or s!="":
                raise PngChunkIncompleteError, "Read too few bytes at %d" % self.tell()
        self.position += len(s)
        return s
    
    def rewind(self):
        """
        Reset the file position to the beginning of the data section.
        
        Note: doesn't redo validation of the magic header, but begins at the
        first chunk
        """
        if self.position < 8:
            raise PngSeekError, "Rewinding before header was read"
        
        try:
            self.file.seek(8)
            self.position = 8
        except:
            raise PngSeekError, "Can't rewind file"
    
    def chunkiter(self):
        """
        This function is an iterator for the chunks in the file.  It basically
        wraps read_chunk and returns a tuple just like read_chunk does.
        """
        while 1:
            try:
                c = self.read_chunk()
            except PngEofError:
                break
            yield PngChunk(*c)
    
    def read_chunk(self):
        """
        Reads a chunk from a png file.  If the file ends prematurely, it will try
        to throw a PngChunkIncompleteError.  If the file ends on a nice boundary
        (and looks like a good file), it will throw PngEofError.  Any other
        exceptions are probably thrown by the underlying file subsystem, and
        should generally be taken to invalidate the object.
        
        Returns a tuple(length(int), type(str*4), data(str*?), and crc(str*4))
        or throws an exception if a whole chunk couldn't be read.
        """
        
        s = self._read(4, failoneof=False)
        if s == "":
            # Surely there's a more Python-like way to do this... I don't like
            # returning "" any more than raising an exception.
            raise PngEofError
        _len = struct.unpack("!L", s)[0]
        
        _type = self._read(4)
        
        _data = self._read(_len)
        
        _crc = self._read(4)
        if self.testcrc:
            newcrc = struct.pack("!L", zlib.crc32(_type + _data))
            if newcrc != _crc:
                raise PngChunkCrcError, "Crc of type+data chunks failed at %d" % self.tell()
        
        # TODO: this should return a new PngChunk class, methinks.  Make the
        # __repr__ of this class show its offset too.  Hmm, needed that earlier.
        # Note that we can't tell() on stdin.
        return (_len, _type, _data, _crc,)