#!/usr/bin/env python

from collections import OrderedDict
from struct import Struct

BM_FLAG_RLE = 8
BM_FLAG_RLE_BIG = 32

HDR = Struct('<4sII')
ENTRY = Struct('<8sBBBBBBI')
RLE = Struct('<I')

class PIGException(Exception):
    pass

class TextureInfo(object):
    def __init__(self, xsize, ysize, flag=0, avg_color=0):
        self.xsize = xsize
        self.ysize = ysize
        self.flag = flag
        self.avg_color = avg_color

class PIG(object):
    def __init__(self):
        self.entries = OrderedDict()
    def add(self, name, animated, frame, raw, info):
        key = name, animated, frame
        if key in self.entries:
            raise PIGException('duplicate entry %r' % name)
        self.entries[key] = raw, info
    def add_from_path(self, name, animated, frame, path, info):
        with open(path, 'rb') as f:
            raw = f.read()
        self.add(name, animated, frame, raw, info)
    def add_entries(self, name, frames=1, xsize=64, ysize=64):
        animated = frames > 1
        for i, frame in enumerate(range(frames)):
            self.add(name, animated, frame,
                     bytearray([i & 0xff] * (xsize * ysize)),
                     TextureInfo(xsize=xsize, ysize=ysize))
    def replace(self, name, animated, frame, raw, info):
        key = name, animated, frame
        if key not in self.entries:
            raise PIGException('no entry %r' % name)
        self.entries[key] = raw, info
    def replace_from_path(self, name, animated, frame, path, info):
        with open(path, 'rb') as f:
            raw = f.read()
        self.replace(name, animated, frame, raw, info)
    def remove(self, name, animated, frame):
        key = name, animated, frame
        if key not in self.entries:
            raise PIGException('no entry %r' % name)
        del self.entries[key]
    def write(self, f):
        count = len(self.entries)
        f.write(HDR.pack(b'PPIG', 2, count))
        offsets = {}
        offset = 0
        for key, (raw, info) in self.entries.items():
            name, animated, frame = key
            offsets[key] = offset
            offset += len(raw)
        for key, (raw, info) in self.entries.items():
            name, animated, frame = key
            if animated:
                frame |= 0x40
            size2 = ((info.xsize & 0x0f00) >> 8) | \
                    ((info.ysize & 0x0f00) >> 4)
            xsize = info.xsize & 0xff
            ysize = info.ysize & 0xff
            flag = info.flag
            avg_color = info.avg_color
            f.write(ENTRY.pack(
                name.encode('latin1'), frame, xsize, ysize, size2, flag,
                avg_color, offsets[key]))
        for raw, info in self.entries.values():
            f.write(raw)
    def write_to_path(self, path):
        with open(path, 'wb') as f:
            self.write(f)

def read(buf):
    magic, version, count = HDR.unpack_from(buf)
    if magic != b'PPIG':
        raise PIGException('File not a Descent PIG file')
    if version != 2:
        raise PIGException('Unknown file version')
    pig = PIG()
    body_offset = HDR.size + ENTRY.size * count
    for i in range(count):
        hdr_offset = HDR.size + ENTRY.size * i
        name, frame, xsize, ysize, size2, flag, avg_color, offset = \
                ENTRY.unpack_from(buf, hdr_offset)
        name = name.split(b'\x00')[0].decode('latin1')
        animated = (frame & 0x40) != 0
        frame &= 0x3f
        info = TextureInfo(
            xsize=xsize | ((size2 & 0x0f) << 8),
            ysize=ysize | ((size2 & 0xf0) << 4),
            flag=flag,
            avg_color=avg_color,
        )
        encflag = flag & (BM_FLAG_RLE | BM_FLAG_RLE_BIG)
        if encflag == 0:
            size = info.xsize * info.ysize
        elif encflag == BM_FLAG_RLE or \
                encflag == (BM_FLAG_RLE | BM_FLAG_RLE_BIG):
            size, = RLE.unpack_from(buf, body_offset + offset)
        else:
            raise PIGException('invalid encoding flag')
        raw = buf[body_offset + offset:body_offset + offset + size]
        pig.add(name, animated, frame, raw, info)
    return pig

def read_from_path(path):
    with open(path, 'rb') as f:
        buf = f.read()
    return read(buf)

def example_usage():
    pig = read_from_path('groupa.pig')
    pig.add_entries('mywall')
    pig.add_entries('myanim', 16)
    pig.write_to_path('new.pig')

if __name__ == '__main__':
    example_usage()