#!/usr/bin/python3
import sys
import png

sprite_width = 32
sprite_height = 32

def process_pixdata(outfile, pixdata, frameindex = 0, pix_w=640, pix_h=400):

    pixmask = 15
    y = pix_h * frameindex
    max_y = y + pix_h - 1

    while y <= max_y:
        x = 0
        max_x = pix_w - 1
        pixline = pixdata[y]
        while x <= max_x:
            px1 = pixline[x+0] & pixmask
            px2 = pixline[x+1] & pixmask
            px3 = pixline[x+2] & pixmask
            px4 = pixline[x+3] & pixmask
            px5 = pixline[x+4] & pixmask
            px6 = pixline[x+5] & pixmask
            px7 = pixline[x+6] & pixmask
            px8 = pixline[x+7] & pixmask
            vmem_word = (px1 << 28) | (px2 << 24) | (px3 << 20) | (px4 << 16) | \
                        (px5 << 12) | (px6 << 8)  | (px7 << 4)  | px8

            outfile.write(vmem_word.to_bytes(4, 'big'))
            x += 8
        y += 1


def process_palette(outfile, palette):
    for r,g,b in palette:
        r4 = r >> 4
        g4 = g >> 4
        b4 = b >> 4
        c12 = r4 << 8 | g4 << 4 | b4

        outfile.write(c12.to_bytes(4, 'big'))


def write_header(outfile):
    magic = b'PIct'
    mode = 1
    outfile.write(magic);
    outfile.write(mode.to_bytes(4, 'big'))


def write_sprite_header(outfile):
    magic = b'SPRT'
    mode = 1
    outfile.write(magic);
    outfile.write(mode.to_bytes(4, 'big'))


def write_pict_file(width, height, px, metadata, outfile):
    print("processing as PICT file...")
    if width != 640:
        print("width must be 640, aborting")
        sys.exit(1)
    pixdata = list(px)
    palette = metadata['palette']

    if len(palette) != 16:
        print("palette must have 16 colors, aborting")
        sys.exit(0)

    with open(outfile,'wb') as f:
        write_header(f)
        process_palette(f, palette)
        process_pixdata(f, pixdata)


def write_sprite_file(width, height, px, metadata, outfile):
    print("processing as SPRT file with {} sprites...".format(height//sprite_height))
    if width != sprite_width:
        print("width must be {}, aborting".format(sprite_width))
        sys.exit(1)
    pixdata = list(px)
    palette = metadata['palette']

    if len(palette) != 16:
        print("palette must have 16 colors instead of {}, aborting".format(len(palette)))
        sys.exit(0)

    with open(outfile,'wb') as f:
        write_sprite_header(f)
        process_pixdata(f, pixdata, pix_w=sprite_width, pix_h=height)


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("Usage: {} <infile> <outfile>".format(sys.argv[0]))
        sys.exit(1)

    infile = sys.argv[1]
    outfile = sys.argv[2]

    r = png.Reader(infile)
    p = r.read()
    width, height, px, metadata = p
    if width == 640:
        write_pict_file(width, height, px, metadata, outfile)
    elif width == sprite_width:
        write_sprite_file(width, height, px, metadata, outfile)
    else:
        print("Can't handle this file, need a width of\n640 or {} pixels with 16 color palette.".format(sprite_width))