#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os, struct, array

K_TITLE = "title"
K_CREDITS = "credits"
K_ID = "id"
K_ALBUM = "album"
K_FILENAME = "filename"
K_DESCRIPTION = "description"

def strip(text):
    '''
    Strip a string by removing not only whitespaces but also all the final 0x0

    @param text: the string to strip
    '''
    pos = text.find(chr(0))
    if pos >= 0:
        text = text[:pos]
    return text.strip()

class Data(object):
    """
    A convenient object where attributes can be added dynamically
    """
    pass

class Pixmap(object):
    """
    This class is the base object manipulated by these file:
    - an image
    - the JPEG data
    - the info about the image

    The image is a Array of bits representing the JPEG data
    The info is a dictionary of key (K_XX), texts
    """
    def __init__(self):
        """
        Constructor
        """
        self.jpeg = array.array('B')
        self.info = {}

    def getJPEGSize(self):
        '''
        @return the size of the image directly from the JPEG data
        '''
        idata = iter(self.jpeg)
        width = None
        height = None
        try:
            B1 = idata.next()
            B2 = idata.next()
            if B1 != 0xFF or B2 != 0xD8:
                return -1, -1
            while True:
                byte = idata.next()
                while byte != 0xFF:
                    byte = idata.next()
                while byte == 0xFF:
                    byte = idata.next()
                if byte >= 0xc0 and byte <= 0xc3:
                    idata.next()
                    idata.next()
                    idata.next()
                    height, width = struct.unpack('>HH',"".join(chr(idata.next()) for b in range(4)) ) #@UnusedVariable
                    break
                else:
                    offset = struct.unpack('>H', chr(idata.next()) + chr(idata.next()))[0] - 2
                    for _ in xrange(offset):
                        idata.next()
        except StopIteration:
            pass
        return (width, height)

class WB1(object):
    """
    This class allows to read wb1, wbd and wbz files
    """
    MAGIC_WB1_NOT_ENCODED = 0xE0FFD8FF
    MAGIC_WB1_ENCODED = 0x42425757
    MAGIC_WBZ = 0x6791AB43
    FILE_MARKER = 0x1082CDE1
    KEY2K = {
        'STR_ImageTitle' :      K_TITLE,
        'title' :               K_TITLE,
        'STR_ImageCredit' :     K_CREDITS,
        'credit' :              K_CREDITS,
        'STR_CollectionTitle' : K_ALBUM,
        'albumTitle' :          K_ALBUM,
        'filename' :            K_FILENAME,
        'FIL_ImageFileName' :   K_FILENAME,
        'id' :                  K_ID,
        }

    @staticmethod
    def read(url, data, buildPixmap = True):
        """
        Read the data and read the image data and info

        @param[in] url the url to the file
        @param[in] data the file data stream
        @param[in] buildPixmap if True build the QPixmap object
        @return a Pixmap image
                None if the magic number of the file does not correspond to the handled types of file
        """
        try:
            pixmap = Pixmap()
            data.seek(0)
            magic = struct.unpack('<I', data.read(4))[0]
            if magic == WB1.MAGIC_WB1_NOT_ENCODED:
                WB1.__decodeWB1(url, data, pixmap, buildPixmap)
            elif magic == WB1.MAGIC_WB1_ENCODED:
                WB1.__decodeWB1_Encoded(url, data, pixmap, buildPixmap)
            elif magic == WB1.MAGIC_WBZ:
                WB1.__decodeWBZ(url, data, pixmap, buildPixmap)
            else:
                return None
            return pixmap
        except:
            return None
    @staticmethod
    def decodeJPEG(data, pixmap, buildPixmap, fileSize = 0):
        """
        Decode the encoded JPEG data

        @param[in] data the file data stream, it must be positioned at the beginning of the encoded JPEG data
        @param[in,out] pixmap the Pixmap instance that will contain the image
        @param[in] buildPixmap if True build the QPixmap object
        @param[in] fileSize the size of the JPEG data. If 0, read them to the end of the file.
        """
        code = data.read(4)
        if code == "0000":
            key = 0xA4
        elif code == "1111":
            key = 0xF2
        else:
            raise Exception("Invalid file format")
        if fileSize:
            pixmap.jpeg.fromfile(data, fileSize)
        else:
            try:
                # read bits until EOF
                while True:
                    pixmap.jpeg.fromfile(data, 0x4000)
            except EOFError:
                pass
        # decode the image (NB '~x' returns a negative integer so we use 'x ^ 0xFF' instead)
        for i in xrange(100):
            pixmap.jpeg[i] = (pixmap.jpeg[i + 100] ^ (0xFF ^ pixmap.jpeg[i])) ^ key
        if buildPixmap:
            # check that the data are valid
            if not pixmap.pixmap.loadFromData(pixmap.jpeg.tostring()):
                raise Exception("Invalid file format")
    @staticmethod
    def readJPEG(data, pixmap, buildPixmap, fileSize = 0):
        """
        Read the JPEG data

        @param[in] data the file data stream, it must be positioned at the beginning of the encoded JPEG data
        @param[in,out] pixmap the Pixmap instance that will contain the image
        @param[in] buildPixmap if True build the QPixmap object
        @param[in] fileSize the size of the JPEG data. If 0, read them to the end of the file.
        """
        if fileSize:
            pixmap.jpeg.fromfile(data, fileSize)
        else:
            try:
                # read bits until EOF
                while True:
                    pixmap.jpeg.fromfile(data, 0x4000)
            except EOFError:
                pass
        if buildPixmap:
            # check that the data are valid
            if not pixmap.pixmap.loadFromData(pixmap.jpeg.tostring()):
                raise Exception("Invalid file format")
    @staticmethod
    def __decodeWB1(url, data, pixmap, buildPixmap):
        """
        Reads the contents of a non encoded WB1 file

        @param[in] url the url to the file
        @param[in] data the stream to the file data positioned just after the 4 bytes of the magic number
        @param[in] buildPixmap if True build the QPixmap object
        @param[out] pixmap the resulting array of bytes
        """
        WB1.readJPEG(data, pixmap, buildPixmap)
        pixmap.info.update(PhotosInfo.parse(url).get(url))
        pixmap.info.update(AlbumInfo.parse(url).get())
    @staticmethod
    def __decodeWB1_Encoded(url, data, pixmap, buildPixmap):
        """
        Reads the contents of an encoded WB1 file

        @param[in] url the url to the file
        @param[in] data the stream to the file data positioned just after the 4 bytes of the magic number
        @param[in] buildPixmap if True build the QPixmap object
        @param[out] pixmap the resulting array of bytes
        """
        WB1.decodeJPEG(data, pixmap, buildPixmap)
        pixmap.info.update(PhotosInfo.parse(url).get(url))
        pixmap.info.update(AlbumInfo.parse(url).get())
    @staticmethod
    def __decodeWBZ(url, data, pixmap, buildPixmap): #@UnusedVariable
        """
        Reads the contents of a WBZ file

        @param[in] url the url to the file
        @param[in] data the stream to the file data positioned just after the 4 bytes of the magic number
        @param[in] buildPixmap if True build the QPixmap object
        @param[out] pixmap the resulting Pixmap
        """
        # JPEG data
        firstFileOffset = struct.unpack('<I', data.read(4))[0]
        data.seek(firstFileOffset)
        includedFileMarker = struct.unpack('<I', data.read(4))[0]
        if includedFileMarker != WB1.FILE_MARKER:
            raise Exception(i18n("Invalid file format"))
        headerSize = struct.unpack('<I', data.read(4))[0]
        data.read(4)
        data.read(256) # fileName
        fileSize = struct.unpack('<I', data.read(4))[0]
        data.read(264)
        code = data.read(4)
        if (code == "WWBB"):
            WB1.decodeJPEG(data, pixmap, buildPixmap, fileSize-8)
        else:
            data.seek(firstFileOffset + headerSize)
            WB1.readJPEG(data, pixmap, buildPixmap, fileSize)
        # description (an INI style file)
        includedFileMarker = struct.unpack('<I', data.read(4))[0]
        if includedFileMarker == WB1.FILE_MARKER:
            headerSize = struct.unpack('<I', data.read(4))[0]
            data.read(4)
            data.read(256) # fileName
            fileSize = struct.unpack('<I', data.read(4))[0]
            data.read(264)
            bytes = data.read(fileSize)
            for line in bytes.splitlines():
                try:
                    key, val = line.split('=',1)
                    k = WB1.KEY2K.get(key, None)
                    if k:
                        pixmap.info[k] = val
                except:
                    pass

class WBC(object):
    """
    This class allows to reab wbc files
    """
    MAGIC_WBC = 0x95FA16AB
    PB_MARKER  = 0xF071CDE2

    @staticmethod
    def read(path, data, buildPixmap = True):
        """
        Read the data and read the image data and info

        @param[in] path the path to the file
        @param[in] data the file data stream
        @param[in] buildPixmap if True build the QPixmap objects
        @return (title, pixmaps) where title is the file title and pixmaps is an array of Pixmap
                (None, None) if the magic number of the file does not correspond to the handled types of file
        """
        try:
            pixmaps = []
            title = ""
            data.seek(0, 2)
            data_size = data.tell()
            data.seek(0)
            magic = struct.unpack('<I', data.read(4))[0]
            if magic == WBC.MAGIC_WBC:
                title = WBC.__decodeWBC(path, data, data_size, pixmaps, buildPixmap)
            else:
                return (None, None)
            return (title, pixmaps)
        except:
                return (None, None)

    @staticmethod
    def __decodeWBC(path, data, data_size, pixmaps, buildPixmap): #@UnusedVariable
        """
        Reads the contents of a non encoded WB1 file

        @param[in] path the path to the file
        @param[in] data the stream to the file data positioned just after the 4 bytes of the magic number
        @param[in] data_size the size of the file
        @param[out] title the file title
        @param[out] pixmaps the resulting array of Pixmap
        @param[in] buildPixmap if True build the QPixmap objects
        @return the title of the collection
        """
        struct.unpack('<I', data.read(4))[0] # firstPBOffset
        data.read(4)
        title = strip(data.read(89))
        data.seek(2196)
        nimages = struct.unpack('<I', data.read(4))[0]
        headers = []
        for i in xrange(nimages): #@UnusedVariable
            hImg = Data()
            hImg.PBOffset = struct.unpack('<I', data.read(4))[0]
            hImg.PBSize = struct.unpack('<I', data.read(4))[0]
            data.read(4)
            hImg.PBAdditionDate = struct.unpack('<I', data.read(4))[0]
            data.read(24)
            headers.append(hImg)
        for h in headers:
            pixmap = Pixmap()
            pixmap.info[K_ALBUM] = title
            # read the header
            data.seek(h.PBOffset)
            PBMarker = struct.unpack('<I', data.read(4))[0]
            headerSize = struct.unpack('<I', data.read(4))[0]
            PBSize = struct.unpack('<I', data.read(4))[0]
            pixmap.info[K_FILENAME] = strip(data.read(256))
            pixmap.info[K_TITLE] = strip(data.read(128))
            pixmap.info[K_DESCRIPTION] = strip(data.read(256))
            pixmap.info[K_CREDITS] = strip(data.read(256))
            data.read(8) # originalFileExtension
            JPGFileSize = struct.unpack('<I', data.read(4))[0]
            data.read(4) # BMPFileSize
            data.read(140)
            data.read(4) # dailyDate
            data.read(4) # additionDate
            data.read(4) # fitToScreen
            pixmap.info[K_ID] = strip(data.read(128))
            data.read(96) # childCategory
            data.read(788) # rootCategory
            # read the image
            if (PBMarker != WBC.PB_MARKER) or (h.PBOffset + PBSize > data_size):
                continue
            data.seek(h.PBOffset + headerSize)
            code = data.read(4)
            if (code == "WWBB"):
                WB1.decodeJPEG(data, pixmap, buildPixmap, JPGFileSize-8)
            else:
                data.seek(h.PBOffset + headerSize)
                WB1.readJPEG(data, pixmap, buildPixmap, JPGFileSize)
            #
            pixmaps.append(pixmap)
        return title



for li in sys.argv[1:]:
    print "Now processing %s..." % li
    # let's process it
    try:
        f = open(unicode(li), 'rb')
    except IOError:
        print "IOError opening %s" % li

    try:
        # let's try first a WB[1DZ] file
        pixmap = WB1.read(li, f, buildPixmap = False)
        if pixmap:
            pixmaps = [ pixmap ]
        else:
            # let's try a WBC file
            title, pixmaps = WBC.read(li, f, buildPixmap = False) #@UnusedVariable
    finally:
        del f
    # now let's save the files
    for pixmap in pixmaps:
    # build the output file name
        fnu = pixmap.info.get(K_TITLE, "image") + ".jpg"
        fn = unicode(fnu, 'latin-1');
        # now save it
        of = open(fn, "wb")
        print "Saved %s" % fnu
        try:
            pixmap.jpeg.tofile(of)
        finally:
            of.close()