# 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(input_filepath, 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 bytes]], one inner list for the columns in
    each row.
    """
    with Image.open(input_filepath) as im:
        # 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)

            sprite_bytes = [sprite.convert('RGBA').tobytes() for sprite in cols]
            assert (len(sprite_bytes) == num_cols)
            rows.append(sprite_bytes)

        return rows


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 sprite_sheet in input_file_paths:
        rows.extend(
            get_sprite_sheet_rows(sprite_sheet, sprite_width, sprite_height))

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

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

        for sprites in rows:
            output.write(ctypes.c_uint16(len(sprites)))
            for sprite_bytes in sprites:
                output.write(sprite_bytes)


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())