#!/usr/bin/env python

import os
import re
import struct
from collections import defaultdict
from PIL import Image
ushrtmax = (1<<16)-1

def encode0(im):
    return ''.join([chr(i) for i in list(im.getdata())])

# greedy RLE
# for each pixel, test which encoding manages to encode most data, then apply
# that encoding and look at the next pixel after the encoded chunk
def encode1(im):
    pixels = im.load()
    w,h = im.size
    result = []
    # these function return a tuple of the compressed string and the amount of
    # pixels compressed
    def rle_comp(x,y):
        # find all pixels after the first one with the same color
        color = pixels[x,y]
        if color == 0xff:
            # this color can't be run length encoded
            return raw_comp(x,y)
        else:
            count = 1
            for x in range(x+1,w):
                if pixels[x,y] == color and count < 255:
                    count += 1
                else:
                    break
            return (struct.pack("<BB", color, count-1), count)
    def raw_comp(x,y):
        # read pixels until finding finding two consecutive ones with the same color
        data = [pixels[x,y]]
        for x in range(x+1,w):
            color = pixels[x,y]
            if color != data[-1] and len(data) < 255:
                data.append(color)
            else:
                break
        return (struct.pack("<BB%dB"%len(data), 0xff, len(data)-1, *data), len(data))
    for y in range(h):
        r = ''
        x = 0
        while x < w:
            rlec, rlel = rle_comp(x, y)
            rawc, rawl = raw_comp(x, y)
            # the message that managed to encode more is chosen 
            if rlel > rawl:
                r += rlec
                x += rlel
            else:
                r += rawc
                x += rawl
        result.append(r)
    return result

def encode23chunk(s,e,pixels,y):
    r = ''
    if pixels[s,y] < 8:
        colors = pixels[s,y]
        count = 1
    else:
        colors = [pixels[s,y]]
        count = 0
    for x in range(s+1,e):
        color = pixels[x,y]
        if count > 0:
            # rle was started
            if color == colors and count < 32:
                # same color again, increase count
                count+=1
            else:
                # either new color or maximum length reached, so write current one
                r+=struct.pack("<B", (colors<<5) | (count-1))
                if color < 7:
                    # new rle color
                    colors = color
                    count = 1
                else:
                    # new non rle color
                    colors = [color]
                    count = 0
        else:
            # non rle was started
            if color < 7 or len(colors) > 31:
                # new rle color, or maximum length reached so write current non rle
                r+=struct.pack("<B", (7<<5) | (len(colors)-1))
                r+=struct.pack("<%dB"%len(colors), *colors)
                if color < 7:
                    colors = color
                    count = 1
                else:
                    colors = [color]
                    count = 0
            else:
                # new non rle color, so append it to current
                colors.append(color)
    # write last color
    if count > 0:
        # write rle
        r+=struct.pack("<B", (colors<<5) | (count-1))
    else:
        # write non rle
        r+=struct.pack("<B", (7<<5) | (len(colors)-1))
        r+=struct.pack("<%dB"%len(colors), *colors)
    return r

# this is like encode3 but a line is not split into 32 pixel chuncks
# the reason for this might just be that format 2 images are always 32 pixel wide
def encode2(im):
    pixels = im.load()
    w,h = im.size
    result = []
    for y in range(h):
        result.append(encode23chunk(0,w,pixels,y))
    return result

# this is like encode2 but limited to only encoding blocks of 32 pixels at a time
def encode3(im):
    pixels = im.load()
    w,h = im.size
    result = []
    for y in range(h):
        res = []
        # encode each row in 32 pixel blocks
        for i in range(w/32):
            res.append(encode23chunk(i*32, (i+1)*32, pixels, y))
        result.append(res)
    return result

fmtencoders = [encode0,encode1,encode2,encode3]

def makedef(indir, outdir):
    infiles = defaultdict(list)
    sig = None
    # sanity checks and fill infiles dict
    for f in os.listdir(indir):
        m = re.match('(\d+)_([a-z0-9_]+)_(\d+)_(\d+)_([A-Za-z0-9_]+)_([0-3]).png', f)
        if not m:
            continue
        t,p,bid,j,fn,fmt = m.groups()
        t,bid,j,fmt = int(t),int(bid),int(j),int(fmt)
        im = Image.open(os.sep.join([indir,f]))
        fw,fh = im.size
        lm,tm,rm,bm = im.getbbox() or (0,0,0,0)
        # format 3 has to have width and lm divisible by 32
        if fmt == 3 and lm%32 != 0:
            # shrink lm to the previous multiple of 32
            lm = (lm/32)*32
        w,h = rm-lm,bm-tm
        if fmt == 3 and w%32 != 0:
            # grow rm to the next multiple of 32
            w = (((w-1)>>5)+1)<<5
            rm = lm+w
        im = im.crop((lm,tm,rm,bm))
        if im.mode != 'P':
            print "input images must have a palette"
            return False
        cursig =(t,p,fw,fh,im.getpalette(),fmt)
        if not sig:
            sig = cursig
        else:
            if sig != cursig:
                print "sigs must match - got:"
                print sig
                print cursig
                return False
        if len(fn) > 9:
            print "filename can't be longer than 9 bytes"
            return False
        data = fmtencoders[fmt](im)
        infiles[bid].append((im,t,p,j,fn,lm,tm,fmt,data))

    if len(infiles) == 0:
        print "no input files detected"
        return False

    # check if j values for all bids are correct and sort them in j order in the process
    for bid in infiles:
        infiles[bid].sort(key=lambda t: t[3])
        for k,(_,_,_,j,_,_,_,_,_) in enumerate(infiles[bid]):
            if k != j:
                print "incorrect j value %d for bid %d should be %d"%(j,bid,k)

    t,p,fw,fh,pal,fmt = cursig
    outname = os.path.join(outdir,p)+".def"
    print "writing to %s"%outname
    outf = open(outname, "w+")

    # write the header
    # full width and height are not used and not the same for all frames
    # in some defs, so just putting the last known value
    outf.write(struct.pack("<IIII", t,fw,fh,len(infiles)))
    # write the palette
    outf.write(struct.pack("768B", *pal))

    # the bid table requires 16 bytes for each bid and 13+4 bytes for each entry
    bidtablesize = 16*len(infiles)+sum(len(l)*(13+4) for l in infiles.values())
    # the position after the bid table is the header plus palette plus bid table size 
    curoffset = 16+768+bidtablesize

    for bid,l in infiles.items():
        # write bid and number of frames
        # the last two values have unknown meaning
        outf.write(struct.pack("<IIII",bid,len(l),0,0))
        # write filenames
        for _,_,_,_,fn,_,_,_,_ in l:
            outf.write(struct.pack("13s", fn+".pcx"))
        # write data offsets
        for im,_,_,_,_,_,_,fmt,data in l:
            outf.write(struct.pack("<I",curoffset))
            w,h = im.size
            # every image occupies size depending on its format plus 32 byte header
            if fmt == 0:
                curoffset += 32+len(data)
            elif fmt == 1:
                # 4*height bytes for lineoffsets
                curoffset += 32+4*h+sum(len(d) for d in data)
            elif fmt == 2:
                # 2*height bytes for lineoffsets plus two unknown bytes
                curoffset += 32+2*h+2+sum(len(d) for d in data)
            elif fmt == 3:
                # width/16 bytes per line as offset header
                curoffset += 32+(w/16)*h+sum(sum([len(e) for e in d]) for d in data)

    for bid,l in infiles.items():
        for im,_,p,j,_,lm,tm,fmt,data in l:
            w,h = im.size
            # size
            # format
            # full width and full height
            # width and height
            # left and top margin
            if fmt == 0:
                s = len(data)
                outf.write(struct.pack("<IIIIIIii",s,fmt,fw,fh,w,h,lm,tm))
                buf = ''.join([chr(i) for i in list(im.getdata())])
                outf.write(buf)
            elif fmt == 1:
                s = 4*h+sum(len(d) for d in data)
                outf.write(struct.pack("<IIIIIIii",s,fmt,fw,fh,w,h,lm,tm))
                lineoffs = []
                acc = 4*h
                for d in data:
                    lineoffs.append(acc)
                    acc += len(d)
                outf.write(struct.pack("<"+"I"*h, *lineoffs))
                for i in data:
                    outf.write(i)
            elif fmt == 2:
                s = 2*h+2+sum(len(d) for d in data)
                outf.write(struct.pack("<IIIIIIii",s,fmt,fw,fh,w,h,lm,tm))
                lineoffs = []
                acc = 0
                for d in data:
                    offs = acc+2*h+2
                    if offs > ushrtmax: 
                        print "exceeding max ushort value: %d"%offs
                        return False
                    lineoffs.append(offs)
                    acc += len(d)
                outf.write(struct.pack("<%dH"%h, *lineoffs))
                outf.write(struct.pack("<BB", 0, 0)) # unknown meaning
                for i in data:
                    outf.write(i)
            elif fmt == 3:
                s = (w/16)*h+sum(sum([len(e) for e in d]) for d in data)
                outf.write(struct.pack("<IIIIIIii",s,fmt,fw,fh,w,h,lm,tm))
                # store the offsets for all 32 pixel blocks
                acc = 0
                lineoffs = []
                for d in data:
                    for e in d:
                        offs = acc+(w/16)*h
                        if offs > ushrtmax:
                            print "exceeding max ushort value: %d"%offs
                            return False
                        lineoffs.append(offs)
                        acc += len(e)
                outf.write(struct.pack("<"+"H"*(w/32)*h, *lineoffs))
                for d in data: # line
                    for e in d: # 32 pixel block
                        outf.write(e)
    return True

if __name__ == '__main__':
    import sys
    if len(sys.argv) != 3:
        print "usage: %s indir outdir"%sys.argv[0]
        exit(1)
    ret = makedef(sys.argv[1], sys.argv[2])
    exit(0 if ret else 1)