#!/usr/bin/env python3
from PIL import Image
import sys
import re
import struct
import zlib
import binascii

def process_rgb(w, h, pix):
    data = bytes()
    for j in range(h):
        for i in range(w):
            r, g, b = pix[i, j]
            c = ((r & 0xF8) << 8) | ((g & 0xFC) << 3) | ((b & 0xF8) >> 3)
            data += struct.pack('>H', c)
    return data


def process_grayscale(w, h, pix):
    data = bytes()
    for j in range(h):
        for i in range(w // 2):
            l1, l2 = pix[i * 2, j], pix[i * 2 + 1, j]
            c = (l1 & 0xF0) | (l2 >> 4)
            data += struct.pack('>B', c)
    return data


def process_image(ifn):
    im = Image.open(ifn)
    w, h = im.size
    print('Opened %s ... %d x %d @ %s' % (ifn, w, h, im.mode))

    if im.mode == 'RGB':
        print('Detected RGB mode')
    elif im.mode == 'L':
        if w % 2 > 0:
            print('PNG file must have width divisible by 2')
            return 3
        print('Detected GRAYSCALE mode')
    else:
        print('Unknown mode:', im.mode)
        return 4

    pix = im.load()

    if im.mode == 'RGB':
        ofn = '%s.toif' % ifn[:-4]
        pixeldata = process_rgb(w, h, pix)
    else:
        ofn = '%s.toig' % ifn[:-4]
        pixeldata = process_grayscale(w, h, pix)
    z = zlib.compressobj(level=9, wbits=10)
    zdata = z.compress(pixeldata) + z.flush()
    zdata = zdata[2:-4] # strip header and checksum

    toif = b''
    if im.mode == 'RGB':
        toif += b'TOIf'
    else:
        toif += b'TOIg'
    toif += struct.pack('<HH', w, h)
    toif += struct.pack('<I', len(zdata))
    toif += zdata
    return toif

# encode vendor name, add length byte and padding to multiple of 4
def vendorencode(vname):
    vbin = vname.encode('utf-8')
    vbin = struct.pack('<B',len(vbin)) + vbin
    vbin += b'\0' * (-len(vbin) & 3)
    return vbin

def encodekey(key):
    if len(key) != 64:
        raise Exception("Wrong key length")
    return binascii.unhexlify(key)

def main():
    if len(sys.argv) < 7:
        print('Usage build_vendorheader key1hex,... m version vendorname vendorimage.png vendorimage.bin')
        return 1

    keys = list(map(encodekey, sys.argv[1].split(',')))
    m = int(sys.argv[2])
    (major,minor) = map(lambda x: int(x),
                        re.search('^(\d+)\.(\d+)$', sys.argv[3]).groups())
    vname = sys.argv[4]
    ifn = sys.argv[5]
    ofn = sys.argv[6]
    if not ifn.endswith('.png'):
        print('Must provide PNG file')
        return 2

    timeout = 0
    vheader = b'TRZV' + struct.pack('<IIBBBB', 0, timeout, major, minor, m, len(keys))
    for k in keys:
        vheader += k
    vheader += vendorencode(vname) + process_image(ifn)
    padding = 65 + (-len(vheader) - 65) & 511
    vheader += b'\0' * padding

    # put in length
    vheader = vheader[0:4] + struct.pack('<I', len(vheader)) + vheader[8:]

    with open(ofn, 'wb') as f:
        f.write(vheader)

main()