#!/usr/bin/env python3

"""
32blit sprite file format

# 12 byte common header

- 8: Type of file identifier ("SPRITEPK", "SPRITERW")
- 2: Total file size

# SPRITEPK header (Packed paletted sprite)

A packed sprite can have a variable number of palette colours up to 255.

Data is packed at the number of bits-per-pixel required to index each palette entry.

IE: a two colour image will be packed with 1bpp

- 2: Sprite width
- 2: Sprite height
- 2: Sprite columns
- 2: Sprite rows
- 1: Format (always paletted)
- 1: Palette size
- X * 4: Pallete entries (where X = number of palette entries)
- Y: Image data

Usage: ./sprite-builder packed input-file.png
"input-file.png" should contain no more than 255 colours!

# SPRITERW header (Raw sprite - RGBA, RGB, P or M)

An unpacked sprite is just a block of raw surface data
in either RGBA (4 bytes), RGB (3 bytes) or P/M (1 bytes) format.

Unpacked sprites are *huge* but useful for streaming into the framebuffer from storage.

The max power of two sprite sizes supported by this file format are-

- 64x64 in RGBA and RGB mode
- 128x128 in P mode

- 2: Sprite width
- 2: Sprite height
- 2: Sprite columns
- 2: Sprite rows
- 1: Format (RGBA, RGB, P or M)
- Y: Image data

Usage: ./sprite-builder raw --format RGB565 input-file.png

"""

import argparse
import math
import sys

from bitstring import BitArray
from construct import this, BitStruct, BitsInteger, Struct, Int8ul, Int16ul, Array, PrefixedArray, Computed, Switch, GreedyBytes, Enum, PaddedString
from PIL import Image, ImagePalette


def compute_bit_length(ctx):
    """Compute the required bit length for image data.

    Uses the count of items in the palette to determine how
    densely we can pack the image data.

    """
    if ctx.format == SpriteFormat.enum.p:
        return (len(ctx.palette.colours) - 1).bit_length()
    elif ctx.format == SpriteFormat.enum.rgba:
        return 32
    elif ctx.format == SpriteFormat.enum.rgb:
        return 24
    else:
        return 0


def compute_data_length(ctx):
    """Compute the required data length for palette based images.

    We need this computation here so we can use `math.ceil` and
    byte-align the result.

    """
    if ctx.format == SpriteFormat.enum.p:
        return math.ceil((ctx.width * ctx.height * ctx.bit_length) / 8)
    else:
        return ctx.width * ctx.height


class SmartEnum(Enum):
    """Fudge to support SpriteFormat.enum.p etc"""
    def __init__(self, *args, **kwargs):
        Enum.__init__(self, *args, **kwargs)

        class x():
            pass
        self.enum = x()
        for e in self.encmapping.keys():
            self.enum.__dict__[e] = e


AssetType = Enum(PaddedString(8, "utf8"), packed='SPRITEPK', raw='SPRITERW')

SpriteFormat = SmartEnum(
    Int8ul,
    rgb=0,
    rgba=1,
    p=2,
    m=3
)

asset_rgba = Struct(
    'r' / Int8ul,
    'g' / Int8ul,
    'b' / Int8ul,
    'a' / Int8ul
)

asset_rgb = Struct(
    'r' / Int8ul,
    'g' / Int8ul,
    'b' / Int8ul
)

asset_packed = Struct(
    'width' / Int16ul,
    'height' / Int16ul,
    'columns' / Int16ul,
    'rows' / Int16ul,
    'format' / SpriteFormat,
    'palette' / Struct(
        'colours' / PrefixedArray(Int8ul, asset_rgba),
    ),
    'bit_length' / Computed(compute_bit_length),
    'data_length' / Computed(compute_data_length),
    'data' / Array(this.data_length, Int8ul)
)

asset_raw = Struct( 
    'width' / Int16ul,
    'height' / Int16ul,
    'columns' / Int16ul,
    'rows' / Int16ul,
    'format' / SpriteFormat,
    'bit_length' / Computed(compute_bit_length),
    'data_length' / Computed(compute_data_length),
    'data' / Switch(this.format, {
        'rgba': Array(this.data_length, asset_rgba),
        'rgb': Array(this.data_length, asset_rgb),
        'p': Array(this.data_length, Int8ul),
        'm': Array(this.data_length, Int8ul)
    })
)

asset = Struct(
    'header' / Struct(
        'type' / AssetType,
        'size' / Int16ul,
    ),
    'data' / GreedyBytes
)

typemap = {
    'packed': asset_packed,
    'raw': asset_raw
}


def unpack_asset(asset_data):
    unpacked = asset.parse(asset_data)
    builder = typemap[unpacked.header.type]
    return builder.parse(unpacked.data)


def build_asset(asset_type, asset_data):
    builder = typemap[asset_type]
    packed = builder.build(asset_data)
    total_size = len(packed) + asset.header.sizeof()
    if total_size > 65535:
        raise RuntimeError("Sprite exceeds maximum filesize: {} > 65535".format(total_size))
    return asset.build(dict(
        header=dict(
            type=asset_type,
            size=total_size,
        ),
        data=packed
    ))


def color(r, g, b, a=None):
    a = 0xff if a is None else a
    return dict(r=r, g=g, b=b, a=a)


def rgba_string(c):
    return "rgba({r}, {g}, {b}, {a})".format(**c)


def format_hex(data):
    return ', '.join(['0x{:02x}'.format(x) for x in data])


def load_palette(image):
    palette_colours = []

    if image.palette is None:
        # Divine the palette from the individual pixels in the image
        # they will be inserted into our target palette in the order
        # found in the palette image
        mode = image.mode
        w, h = image.size

        if mode in ['RGB', 'RGBA']:
            for y in range(h):
                for x in range(w):
                    # The spread operator will take both RGB and RGBA
                    # tuples and pass them into color() properly.
                    c = color(*image.getpixel((x, y)))
                    if c not in palette_colours:
                        palette_colours.append(c)

        else:
            raise RuntimeError(f"Unsupported image mode: {mode}")

    else:
        mode, palette = image.palette.getdata()
        palette = list(palette)
        if mode == 'RGB':
            for x in range(len(palette) // 3):
                r, g, b = palette[x * 3:(x * 3) + 3]
                palette_colours.append(color(r, g, b))

        elif mode == 'RGBA':
            for x in range(len(palette) // 4):
                r, g, b, a = palette[x * 4:(x * 4) + 4]
                palette_colours.append(color(r, g, b, a))

    if len(palette_colours) > 255:
        raise RuntimeError(f"Too many colours in palette! ({len(palette_colours)})")

    return palette_colours


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--nowarn', action='store_true', help='Suppress deprecation warning')
    subparsers = parser.add_subparsers(help='Commands', dest='command')
    parser_packed = subparsers.add_parser('packed', help='Process an image file into a paletted 32blit sprite')
    parser_raw = subparsers.add_parser('raw', help='Process an image file into a raw (rgba, rgb) 32blit sprite')

    parser.add_argument('--out', type=str, help='set output filename (default <inputfile>.blit)', default=None)
    parser.add_argument('file', type=str, help='input file')

    # The engine presently only supports 8x8 tile sizes
    parser.add_argument('--tilesizex', type=int, help='tile size x', default=8)
    parser.add_argument('--tilesizey', type=int, help='tile size y', default=8)

    # Allow for naming of the resulting packed_data array
    parser_packed.add_argument('--arrayname', type=str, help='set output array name (default <packed_data>)', default='packed_data')

    # Raw input can be given as any image type, and will be packed into:
    # * RGBA (full fidelity)
    # * RGB888 (discard alpha)
    # * RGB565 (discard alpha, convert colours with bitshift: R >> 3, G >> 2 B >> 3)
    parser_raw.add_argument('--format', type=str, choices=('RGBA', 'RGB'), help='raw data output format', default='RGBA')

    # A palette file should be a PNG file containing a series of pixels in the desired colours and order of the output palette.
    # If it's indexed, 8bpp colour then its palette will be used directly, otherwise the pixel data will be converted automatically.
    parser_packed.add_argument('--palette', type=str, help='image or palette file', default=None)

    args = parser.parse_args()

    if not args.nowarn:
        print('''
This tool has been deprecated in favour of https://github.com/pimoroni/32blit-tools
Run with --nowarn to hide this message.
''')

    source_image = Image.open(args.file)

    if args.command == 'raw':
        # Make sure we're always working with the same type of image
        source_image = source_image.convert('RGBA')

        raw_data = []

        for y in range(source_image.height):
            for x in range(source_image.width):
                r, g, b, a = source_image.getpixel((x, y))
                if args.format == 'RGBA':
                    raw_data.append({'r': r, 'g': g, 'b': b, 'a': a})
                if args.format == 'RGB':
                    raw_data.append({'r': r, 'g': g, 'b': b})

        raw = build_asset(
            'raw',
            dict(
                width=source_image.width,
                height=source_image.height,
                columns=source_image.width // args.tilesizex,
                rows=source_image.height // args.tilesizey,
                format=args.format.lower(),
                data=raw_data
            )
        )

        if args.out is not None:
            with open(args.out, 'wb+') as f:
                f.write(raw)
            sys.exit(0)

    if args.command == 'packed':
        target_image = source_image
        source_mode = source_image.mode
        source_palette_colours = load_palette(source_image)
        target_palette_colours = source_palette_colours

        bit_length = (len(source_palette_colours) - 1).bit_length()

        if source_mode == 'P':
            if args.palette is not None:
                palette_image = Image.open(args.palette)
                target_palette_colours = load_palette(palette_image)

                w, h = source_image.size
                for y in range(h):
                    for x in range(w):
                        p = source_image.getpixel((x, y))
                        c = source_palette_colours[p]
                        try:
                            new_p = target_palette_colours.index(c)
                            target_image.putpixel((x, y), new_p)
                        except ValueError:
                            raise RuntimeError("Colour {},{},{},{} not found in target palette!".format(*c.rgba()))

                source_palette_colours = target_palette_colours
        else:
            target_image = Image.new('P', source_image.size)
            w, h = source_image.size
            for y in range(h):
                for x in range(w):
                    if source_mode == 'RGBA':
                        r, g, b, a = source_image.getpixel((x, y))
                        c = color(r, g, b, a)
                    elif source_mode == 'RGB':
                        r, g, b = source_image.getpixel((x, y))
                        c = color(r, g, b)

                    try:
                        new_p = target_palette_colours.index(c)
                        target_image.putpixel((x, y), new_p)
                    except ValueError:
                        raise RuntimeError("Colour {},{},{},{} not found in target palette!".format(*c.rgba()))


        packed_image_data = BitArray().join(BitArray(uint=x, length=bit_length) for x in target_image.tobytes()).tobytes()

        packed = build_asset(
            'packed',
            dict(
                width=target_image.width,
                height=target_image.height,
                columns=target_image.width // args.tilesizex,
                rows=target_image.height // args.tilesizey,
                format='p',
                palette=dict(
                    colours=source_palette_colours
                ),
                data=packed_image_data
            )
        )

        if args.out is not None:
            with open(args.out, 'wb+') as f:
                f.write(packed)
            sys.exit(0)

        packed = list(packed)
        formatted_data = []
        formatted_palette = []
        row_width = 16

        HEADER_SIZE = 20

        PALETTE_SIZE = len(source_palette_colours) * 4

        palette_data = packed[HEADER_SIZE:HEADER_SIZE + PALETTE_SIZE]
        for x in range(0, math.ceil(len(palette_data) / 4)):
            offset = x * 4
            current_row = palette_data[offset:offset + 4]
            formatted_palette.append(format_hex(current_row))

        packed_data = packed[HEADER_SIZE + PALETTE_SIZE:]
        for x in range(0, math.ceil(len(packed_data) / float(row_width))):
            offset = x * row_width
            current_row = packed_data[offset:offset + row_width]
            formatted_data.append(format_hex(current_row))

        print("""
uint8_t {array_name}[] = {{
    {type}, // type: spritepk (packed, paletted sprite)
    {payload_size}, // payload size ({comment_payload_size})

    {width}, // width ({comment_width})
    {height}, // height ({comment_height})
    {cols}, // cols ({comment_cols})
    {rows}, // rows ({comment_rows})

    {format}, // format (paletted)

    {colours}, // number of palette colours ({number_of_colours})
//  r     g     b     a
    {palette},

    {data}
}};
""".format(
            array_name=args.arrayname,
            type=format_hex(packed[0:8]),
            payload_size=format_hex(packed[8:10]),
            comment_payload_size=len(packed),
            width=format_hex(packed[10:12]),
            comment_width=target_image.width,
            height=format_hex(packed[12:14]),
            comment_height=target_image.height,
            cols=format_hex(packed[14:16]),
            comment_cols=target_image.width // args.tilesizex,
            rows=format_hex(packed[16:18]),
            comment_rows=target_image.height // args.tilesizey,
            format=format_hex(packed[18:19]),
            colours=format_hex(packed[19:20]),
            number_of_colours=packed[19:20][0],
            data=',\n    '.join(formatted_data),
            palette=',\n    '.join(formatted_palette)
        ))