# Converts assets to binary formats (.ts, .tm, .ss) for the engine.
#
# Input file formats:
#  - Tiled tile set (.tsx)
#  - Tiled tile map (.tmx)
#  - Sprite sheets (.jpg, .png, etc), 1 row per animation.
#
# Output file formats:
#  - Binary tile set file (.ts)
#  - Binary tile map file (.tm)
#  - Binary sprite sheet file (.ss)
#
import argparse
import ctypes
import os
from PIL import Image
import sys
from xml.etree import ElementTree

# Maximum length of path strings in .TS and .TM files.
# Must match the engine's value.
MAX_PATH_LENGTH = 128


def drop_extension(filepath):
    return filepath[:filepath.rfind('.')]


def to_char_array(string, length):
    """Convert a string to a fixed-length ASCII char array.

    The length of str must be at most length-1 so that the resulting string can
    be null-terminated.
    """
    assert (len(string) < length)
    chars = string.encode("ascii")
    nulls = ("\0" * (length - len(string))).encode("ascii")
    return chars + nulls


def convert_tsx(input_filepath, output_filepath):
    """Converts a Tiled .tsx tileset file to a .TS tile set file."""
    xml = ElementTree.parse(input_filepath)
    root = xml.getroot()

    tile_count = int(root.attrib["tilecount"])
    max_tile_width = int(root.attrib["tilewidth"])
    max_tile_height = int(root.attrib["tileheight"])

    print(f"Tile count: {tile_count}")
    print(f"Max width:  {max_tile_width}")
    print(f"Max height: {max_tile_height}")

    with open(output_filepath, 'bw') as output:
        output.write(ctypes.c_uint16(tile_count))
        output.write(ctypes.c_uint16(max_tile_width))
        output.write(ctypes.c_uint16(max_tile_height))

        num_tile = 0
        for tile in root:
            # Skip the "grid" and other non-tile elements.
            if not tile.tag == "tile":
                continue

            # Assuming tiles are numbered 0..N.
            tile_id = int(tile.attrib["id"])
            assert (tile_id == num_tile)
            num_tile += 1

            image = tile[0]
            tile_width = int(image.attrib["width"])
            tile_height = int(image.attrib["height"])
            tile_path = image.attrib["source"]

            output.write(ctypes.c_uint16(tile_width))
            output.write(ctypes.c_uint16(tile_height))

            with Image.open(tile_path) as im:
                bytes = im.convert('RGBA').tobytes()
                output.write(bytes)


def convert_tmx(input_filepath, output_filepath):
    """Converts a Tiled .tmx file to a .TM tile map file."""
    xml = ElementTree.parse(input_filepath)
    root = xml.getroot()

    map_width = int(root.attrib["width"])
    map_height = int(root.attrib["height"])
    base_tile_width = int(root.attrib["tilewidth"])
    base_tile_height = int(root.attrib["tileheight"])
    num_layers = 1

    print(f"Map width:   {map_width}")
    print(f"Map height:  {map_height}")
    print(f"Tile width:  {base_tile_width}")
    print(f"Tile height: {base_tile_height}")

    with open(output_filepath, 'bw') as output:
        output.write(ctypes.c_uint16(map_width))
        output.write(ctypes.c_uint16(map_height))
        output.write(ctypes.c_uint16(base_tile_width))
        output.write(ctypes.c_uint16(base_tile_height))
        output.write(ctypes.c_uint16(num_layers))

        tileset_path = None

        for child in root:
            if child.tag == "tileset":
                tileset = child
                tileset_path = tileset.attrib["source"]

                print(f"Tile set: {tileset_path}")

                tileset_path = tileset_path.replace("tsx", "ts")
            elif child.tag == "layer":
                layer = child
                layer_id = int(layer.attrib["id"])
                layer_width = int(layer.attrib["width"])
                layer_height = int(layer.attrib["height"])

                print(f"Layer:  {layer_id}")
                print(f"Width:  {layer_width}")
                print(f"Height: {layer_height}")

                assert (tileset_path)
                output.write(to_char_array(tileset_path, MAX_PATH_LENGTH))

                # Assume the layer's dimensions matches the map's.
                assert (layer_width == map_width)
                assert (layer_height == map_height)

                data = layer[0]
                # Handle other encodings later.
                assert (data.attrib["encoding"] == "csv")

                csv = data.text.strip()
                rows = csv.split('\n')
                for row in rows:
                    tile_ids = [x.strip() for x in row.split(',') if x]
                    for tile_id in tile_ids:
                        output.write(ctypes.c_uint16(int(tile_id)))


def get_num_cols(image, sprite_width):
    """Return the number of non-empty columns in the image.

    Assumes no gaps in the columns.
    """
    assert (image.width % sprite_width == 0)
    num_cols = image.width // sprite_width

    # Start the search from right to left.
    for col in reversed(range(1, num_cols)):
        left = (col - 1) * sprite_width
        right = col * sprite_width
        rect = image.crop((left, 0, right, image.height))
        min_max = rect.getextrema()
        for (channel_min, channel_max) in min_max:
            if channel_min != 0 or channel_max != 0:
                # 'col' is the rightmost non-empty column.
                # Assuming no gaps, col+1 is the number of non-empty columns.
                return col + 1

    return 0


def get_sprite_sheet_rows(im, sprite_width, sprite_height):
    """Gets the individual rows of a sprite sheet.

    The input sprite sheet can have any number of rows.

    Returns a list of lists [[sprite]], one inner list for the columns in each
    row.
    """
    # Sprite sheet's width and height must be integer multiples of the
    # sprite's width and height.
    assert (im.width % sprite_width == 0)
    assert (im.height % sprite_height == 0)

    num_rows = im.height // sprite_height

    rows = []
    for row in range(num_rows):
        # Get the number of columns.
        upper = row * sprite_height
        lower = (row + 1) * sprite_height
        whole_row = im.crop((0, upper, im.width, lower))
        num_cols = get_num_cols(whole_row, sprite_width)
        assert (num_cols > 0)

        # Crop the row into N columns.
        cols = []
        for i in range(num_cols):
            left = i * sprite_width
            right = (i + 1) * sprite_width
            sprite = im.crop((left, upper, right, lower))
            cols.append(sprite)

        assert (len(cols) == num_cols)
        rows.append(cols)

    return rows


def make_image_from_rows(rows, sprite_width, sprite_height):
    """Concatenate the rows into a single RGBA image."""
    im_width = sprite_width * max(len(row) for row in rows)
    im_height = len(rows) * sprite_height
    im = Image.new('RGBA', (im_width, im_height))
    y = 0
    for row in rows:
        x = 0
        for sprite in row:
            im.paste(sprite.convert('RGBA'), (x, y))
            x += sprite_width
        y += sprite_height
    return im


def convert_sprite_sheet(input_file_paths, sprite_width, sprite_height,
                         output_filepath):
    """Converts a set of sprite sheet images into a binary sprite sheet file
    (.ss).

    The input sprite sheets can have any number of rows, one row per animation.
    All rows from all sprite sheets are concatenated in the output file.

    The sprite's width and height is assumed constant throughout the input
    sprite sheets.
    """
    rows = []
    for input_filepath in input_file_paths:
        with Image.open(input_filepath) as sprite_sheet:
            rows.extend(
                get_sprite_sheet_rows(sprite_sheet, sprite_width,
                                      sprite_height))

    im = make_image_from_rows(rows, sprite_width, sprite_height)
    im = im.convert(mode="P", palette=Image.ADAPTIVE, colors=256)

    # The sprite data in 'rows' is no longer needed.
    # Keep just the number of columns per row.
    rows = [len(row) for row in rows]

    with open(output_filepath, 'bw') as output:
        output.write(ctypes.c_uint16(sprite_width))
        output.write(ctypes.c_uint16(sprite_height))
        output.write(ctypes.c_uint16(len(rows)))

        # Write palette.
        # getpalette() returns 256 colors, but the palette might use less than
        # that. getcolors() returns the number of unique colors.
        # getpalette() also returns a flattened list, which is why we must *4.
        num_colours = len(im.getcolors())
        colours = im.getpalette(rawmode="RGBA")[:4 * num_colours]
        palette = []
        for i in range(0, 4 * num_colours, 4):
            palette.append((colours[i], colours[i + 1], colours[i + 2],
                            colours[i + 3]))

        output.write(ctypes.c_uint16(len(palette)))
        output.write(bytearray(colours))

        print(f"Sprite width:  {sprite_width}")
        print(f"Sprite height: {sprite_height}")
        print(f"Rows:          {len(rows)}")
        print(f"Colours:       {len(palette)}")

        # print("Palette")
        # for i, colour in enumerate(palette):
        #     print(f"{i}: {colour}")

        for row, num_columns in enumerate(rows):
            output.write(ctypes.c_uint16(num_columns))
            upper = row * sprite_height
            lower = (row + 1) * sprite_height
            for col in range(num_columns):
                left = col * sprite_width
                right = (col + 1) * sprite_width
                sprite = im.crop((left, upper, right, lower))
                sprite_bytes = sprite.tobytes()

                assert (len(sprite_bytes) == sprite_width * sprite_height)
                output.write(sprite_bytes)

                # if (row == 0) and (col == 0):
                #     print(f"Sprite: ({len(sprite_bytes)})")
                #     print(list(sprite_bytes))
                #     sprite.save("out.png")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input",
                        nargs="+",
                        help="Input file (.tsx, .tmx) or path regex (sprite sheets)")
    parser.add_argument("--width", type=int, help="Sprite width in pixels")
    parser.add_argument("--height", type=int, help="Sprite height in pixels")
    parser.add_argument("--out", help="Output file (sprite sheets)")
    args = parser.parse_args()

    if ".tsx" in args.input:
        output_filepath_no_ext = drop_extension(args.input)
        output_filepath = output_filepath_no_ext + ".ts"
        convert_tsx(args.input, output_filepath)
    elif ".tmx" in args.input:
        output_filepath_no_ext = drop_extension(args.input)
        output_filepath = output_filepath_no_ext + ".tm"
        convert_tmx(args.input, output_filepath)
    else:
        # Sprite sheets.
        if not args.width or not args.height:
            print("Sprite width and height must be given")
            return 1
        output_filepath = args.out if args.out else "out.ss"
        convert_sprite_sheet(args.input, args.width, args.height,
                             output_filepath)

    return 0


if __name__ == '__main__':
    sys.exit(main())