#!/usr/bin/env python3

"""
32blit asset file format (extension always .blit)

12 byte common header

- 8: Type of file identifier ("SPRITE  ", "MAP     ", "SONG    ", "SFX     ")
- 4: Total file size

SPRITE header

- 2: Sprite width
- 2: Sprite height
- 2: Sprite columns
- 2: Sprite rows
- 1: Format (RGBA, RGB, 8-bit palette, 6-bit palette + 2-bit alpha)
- 1: Palette size (if relevant)
- 1: Palette transparent colour index
- X * 4: Pallete entries (where X = number of palette entries)
- Y: Image data

Unpaletted - RGB, RGBA
Paletted - packed 2bpp (4 colours), 3bpp (8 colours), 4bpp (16 colours), etc
"""

import argparse
import math
import sys

from bitstring import BitArray
from construct import this, 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 8 * 4
    elif ctx.format == SpriteFormat.enum.rgb:
        return 8 * 3
    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"), sprite='SPRITE', map='MAP', song='SONG', sfx='SFX')

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

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

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

asset_sprite = 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' / Switch(this.format, {
        'rgba': Array(this.data_length, asset_rgba),
        'rgb': Array(this.data_length, asset_rgb),
        'p': Array(this.data_length, Int8ul)
    })
)

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

typemap = {
    'sprite': asset_sprite
}


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)
    return asset.build(dict(
        header=dict(
            type=asset_type,
            size=len(packed) + asset.header.sizeof(),
        ),
        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)


"""
class color():

    def __init__(self, r, g, b, a=None):
        a = 0xff if a is None else a
        self.r = r
        self.g = g
        self.b = b
        self.a = a

    def rgba(self):
        return (self.r, self.g, self.b, self.a)

    def rgba_dict(self):
        return {'r': self.r, 'g': self.g, 'b': self.b, 'a': self.a}

    def __eq__(self, other):
        return self.rgba() == other.rgba()

    def __iter__(self):
        return self.rgba_dict().__iter__()
"""


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_sprite = subparsers.add_parser('sprite', help='Process an image file into a 32blit asset')
    parser_map = subparsers.add_parser('map', help='Process a map file into a 32blit asset')

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

    # TODO: Right now unpacked data is ridiculously huge and probably impractical to use,
    # either we need a hard limit on the size of assets built in RGBA or RGB native formats (24/32bpp)
    # or drop output format altogether.
    parser_sprite.add_argument('--format', type=str, choices=('RGB', 'RGBA', 'P'), help='data output format (UNIMPLEMENTED)')

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

    # 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_sprite.add_argument('--palette', type=str, help='image or palette file', default=None)

    parser_map.add_argument('--sizex', type=int, help='map size x')
    parser_map.add_argument('--sizey', type=int, help='map size y')

    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.
''')

    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 args.command == 'sprite':

        source_image = Image.open(args.file)
        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(
            'sprite',
            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)

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

        # formatted_data = []
        # row_width = image.width // 16
        # for x in range(0, len(packed_image_data) // row_width):
        #     offset = x * row_width
        #     current_row = packed_image_data[offset:offset + row_width]
        #     formatted_data.append(', '.join(['0x{:02x}'.format(x) for x in current_row]))

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

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

        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 packed_data[] = {{
    {type}, // type: 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

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

    {data}
}};
""".format(
            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]),
            data=',\n    '.join(formatted_data),
            palette=',\n    '.join(formatted_palette)
        ))


#         print("""
# rgba palette_data[] = {{
#     {palette_data}
# }};
#
# uint8_t sprite_data[] = {{
#     {image_data}
# }};
# """.format(
#     palette_data=',\n    '.join([rgba_string(x) for x in palette_colours]),
#     image_data=',\n    '.join(formatted_data)
# ))