"""
Mac OS .icns file decoder.  If you're using PIL, just import this module and
it will be registered as a codec.  Otherwise, use IcnsFile to get the data.

Note that it currently only decodes 32bit RGB icons and 8bit masks!
"""

import struct
from array import array
from itertools import *

__all__ = ['IcnsFile']

HEADERSIZE = 8
def nextheader(fobj):
    return struct.unpack('>4sI', fobj.read(HEADERSIZE))

def arrayfromfile(code, fobj, size):
    a = array(code)
    a.fromfile(fobj, size)
    return a

def read_32t(fobj, (start, length), (width, height)):
    # The 128x128 icon seems to have an extra header for some reason.
    fobj.seek(start)
    sig = fobj.read(4)
    if sig != '\x00\x00\x00\x00':
        raise SyntaxError, 'Unknown signature, expecting 0x00000000'
    return read_32(fobj, (start + 4, length - 4), (width, height))

def read_32(fobj, (start, length), (width, height)):
    """
    Read a 32bit RGB icon resource.  Seems to be either uncompressed or
    an RLE packbits-like scheme
    """
    fobj.seek(start)
    sizesq = width * height
    indata = arrayfromfile('B', fobj, length)
    if length == sizesq * 3:
        # uncompressed
        return {
            'R': indata[0::3],
            'G': indata[1::3],
            'B': indata[2::3],
        }
    channels = {}
    iterdata = iter(indata)
    for channel in 'RGB':
        data = channels[channel] = array('B')
        bytesleft = sizesq
        append = data.append
        for byte in iterdata:
            if byte & 0x80:
                blocksize = byte - 125
                ext = repeat(iterdata.next(), blocksize)
            else:
                blocksize = byte + 1
                ext = islice(iterdata, 0, blocksize)
            for byte in ext:
                append(byte)
            bytesleft -= blocksize
            if bytesleft <= 0:
                break
        if bytesleft != 0:
            raise SyntaxError, "Error reading %r channel [%r]" % (channel, bytesleft)
    if list(iterdata):
        raise SyntaxError, "Extra data!"
    return channels

def read_mk(fobj, (start, length), (width, height)):
    # Alpha masks seem to be uncompressed
    fobj.seek(start)
    return {
        'A': arrayfromfile('B', fobj, width*height),
    }

class IcnsFile(object):
    """
    This only reads the 32bit icon resources and the 8 bit masks.
    
    Anything else, would be uncivilized.
    """
    SIZES = {
        (128, 128): [
            ('it32', read_32t),
            ('t8mk', read_mk),
        ],
        (48, 48): [
            ('ih32', read_32),
            ('h8mk', read_mk),
        ],
        (32, 32): [
            ('il32', read_32),
            ('l8mk', read_mk),
        ],
        (16, 16): [
            ('is32', read_32),
            ('s8mk', read_mk),
        ],
    }
    
    def __init__(self, fobj):
        """
        fobj is a file-like object as an icns resource
        """
        # signature : (start, length)
        self.dct = dct = {}
        self.fobj = fobj
        sig, filesize = nextheader(fobj)
        if sig != 'icns':
            raise SyntaxError, 'not an icns file'
        i = HEADERSIZE
        while i < filesize:
            sig, blocksize = nextheader(fobj)
            i += HEADERSIZE
            blocksize -= HEADERSIZE
            dct[sig] = (i, blocksize)
            fobj.seek(blocksize, 1)
            i += blocksize

    def itersizes(self):
        for size, fmts in type(self).SIZES.iteritems():
            for (fmt, reader) in fmts:
                if fmt in self.dct:
                    yield size
                    break
    
    def bestsize(self):
        sizes = list(self.itersizes())
        if not sizes:
            raise SyntaxError, "No 32bit icon resources found"
        return max(sizes)

    def dataforsize(self, size):
        """
        Get an icon resource as {channel: array}.  Note that
        the arrays are bottom-up like windows bitmaps and will likely
        need to be flipped or transposed in some way.
        """
        dct = {}
        for code, reader in type(self).SIZES[size]:
            desc = self.dct.get(code)
            if desc is not None:
                dct.update(reader(self.fobj, desc, size))
        return dct

#
# PIL support
#

try:
    import Image
    from ImageFile import ImageFile
except ImportError:
    #
    # Bogus stuff so that it imports without PIL installed
    #
    class Ignore(object):
        def __getattr__(self, *args, **kwargs):
            return self
        def __call__(self, *args, **kwargs):
            return self
    Image = Ignore()
    ImageFile = object

def to_image(icns, size=None):
    if size is None:
        size = icns.bestsize()
    channels = icns.dataforsize(size)
    bands = [Image.frombuffer('L', size, channels[mode]) for mode in 'RGBA']
    return Image.merge('RGBA', bands).transpose(Image.FLIP_TOP_BOTTOM)

class IcnsImageFile(ImageFile):
    """
    PIL read-only image support for Mac OS .icns files.
    Chooses the best resolution, but will possibly load
    a different size image if you mutate the size attribute.

    The info dictionary has a key 'sizes' that is a list
    of sizes that the icns file has.
    """
    format = "icns"
    format_description = "Mac OS icns resource"

    def _open(self):
        self.icns = IcnsFile(self.fp)
        self.mode = 'RGBA'
        self.size = self.icns.bestsize()
        self.info['sizes'] = list(self.icns.itersizes())
        # Just use this to see if it's loaded or not yet.
        self.tile = ('',)

    def load(self):
        Image.Image.load(self)
        if not self.tile:
            return
        self.load_prepare()
        # This is likely NOT the best way to do it, but whatever.
        self.im = to_image(self.icns, self.size).im
        self.fp = None
        self.icns = None
        self.tile = ()
        self.load_end()

Image.register_open('icns', IcnsImageFile)
Image.register_extension('icns', '.icns')

if __name__ == '__main__':
    import sys
    icns = IcnsFile(sys.argv[1])
    for size in icns.itersizes():
        icns.dataforsize(size)