# 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 sys
from typing import Generator
from xml.etree import ElementTree

from PIL import Image

# 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) -> bytes:
    """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 load_image(path) -> bytes:
    """Load an image as an array of RGBA bytes."""
    with Image.open(path) as im:
        return im.convert('RGBA').tobytes()


def carve_image(rgba_bytes, tile_width, tile_height, columns) -> Generator[bytearray]:
    """Carve out individual tiles from a single-image tile set."""
    # Image dimensions in pixels.
    image_width = columns * tile_width
    image_height = len(rgba_bytes) // image_width // 4

    tile_bytes = bytearray(tile_width * tile_height * 4)
    for image_y0 in range(0, image_height, tile_height):  # y-origin of tile inside image
        for image_x0 in range(0, image_width, tile_width):  # x-origin of tile inside image
            for y in range(tile_height):
                image_y = image_y0 + y  # y of current pixel inside image
                for x in range(tile_width):
                    image_x = image_x0 + x  # x of current pixel inside image
                    tile_bytes[(y * tile_width + x) * 4] = (
                        rgba_bytes)[(image_y * image_width + image_x) * 4]
            yield tile_bytes.copy()


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

    # Header.
    tileset_tile_count = int(tileset.attrib["tilecount"])
    tileset_tile_width = int(tileset.attrib["tilewidth"])
    tileset_tile_height = int(tileset.attrib["tileheight"])

    print(f"Tile count:  {tileset_tile_count}")
    print(f"Tile width:  {tileset_tile_width}")
    print(f"Tile height: {tileset_tile_height}")

    pixels = []  # List of byte arrays
    pixels_offset = 0

    def output_tile(output, tile_width, tile_height, tile_image_bytes):
        # Expecting RGBA pixels.
        assert (len(tile_image_bytes) == (tile_width * tile_height * 4))

        nonlocal pixels_offset
        tile_pixels_offset = pixels_offset
        pixels.append(tile_image_bytes)
        pixels_offset += len(tile_image_bytes)

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

    with open(output_filepath, 'bw') as output:
        # Write the header.
        output.write(ctypes.c_uint16(tileset_tile_count))
        output.write(ctypes.c_uint16(tileset_tile_width))
        output.write(ctypes.c_uint16(tileset_tile_height))
        output.write(ctypes.c_uint16(0))  # Pad.

        # A tileset made up of multiple images contains various 'tile' children,
        # each with their own 'image':
        #
        #   <tile id="0">
        #     <image width="32" height="32" source="separated images/tile_000.png"/>
        #   </tile>
        #   <tile id="1">
        #     <image width="32" height="32" source="separated images/tile_001.png"/>
        #   </tile>
        #
        # A tileset made up of a single image contains a single 'image' child:
        #
        #   <image source="tileset.png" width="416" height="368"/>
        num_tile = 0
        for child in tileset:
            if child.tag == "image":
                # This is a single-image tileset.
                image = child
                image_path = image.attrib["source"]
                image_bytes = load_image(image_path)

                # We expect the 'columns' attribute to be >0 for a single-image
                # tile set.
                columns = int(tileset.attrib["columns"])
                assert (columns > 0)

                # Tile dimensions are those given by the tileset's baseline width and height.
                tile_width = tileset_tile_width
                tile_height = tileset_tile_height

                # Cut each of the WxH tiles, and store them in the output in "tile order".
                for tile_bytes in carve_image(image_bytes, tile_width, tile_height, columns):
                    output_tile(output, tile_width, tile_height, tile_bytes)

                # Make sure not to process 'tile' elements, which may exist in
                # a single-image tileset to define probabilities, for example.
                break

            elif child.tag == "tile":
                # This is a tile image in a tileset made of a collection of images.
                tile = child

                # We assume that tiles are numbered 0..N-1. Assert this.
                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"])
                image_path = image.attrib["source"]

                # Tile dimensions must be a multiple of the tileset's baseline
                # tile width and height.
                assert (tile_width > 0)
                assert (tile_height > 0)
                assert ((tile_width % tileset_tile_width) == 0)
                assert ((tile_height % tileset_tile_height) == 0)

                image_bytes = load_image(image_path)
                output_tile(output, tile_width, tile_height, image_bytes)

        # Write the pixel data.
        for bytes in pixels:
            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}")

    tileset_path = None

    with open(output_filepath, 'bw') as output:
        for child in root:
            if child.tag == "tileset":
                assert (not tileset_path)  # Only supporting one tile set per map right now.

                tileset = child
                tileset_path = tileset.attrib["source"]

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

                tileset_path = tileset_path.replace("tsx", "ts")

                # Write the header.
                output.write(to_char_array(tileset_path, MAX_PATH_LENGTH))
                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))
            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}")

                # 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:
                        # TODO: We need to handle 'firsgid' in TMX files.
                        #  For now, assume it's 1 and do -1 to make the tiles 0-based.
                        output.write(ctypes.c_uint16(int(tile_id) - 1))


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())
        palette = bytearray(im.getpalette(rawmode="RGBA")[:4 * num_colours])
        assert (num_colours == (len(palette) // 4))

        output.write(ctypes.c_uint16(num_colours))
        output.write(palette)

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

        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():
    # TODO: Use a subparser for each type of input file for clarity of arguments.
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    # Tile sets.
    tileset_parser = subparsers.add_parser("tileset")
    tileset_parser.add_argument("input", help="Input file (.tsx)")
    # Tile maps.
    tilemap_parser = subparsers.add_parser("tilemap")
    tilemap_parser.add_argument("input", help="Input file (.tmx)")
    # Sprite sheet.
    sprite_parser = subparsers.add_parser("sprite")
    sprite_parser.add_argument("input", nargs="+", help="Input files")
    sprite_parser.add_argument("-W", "--width", type=int, required=True, help="Sprite width in pixels")
    sprite_parser.add_argument("-H", "--height", type=int, required=True, help="Sprite height in pixels")
    sprite_parser.add_argument("-o", "--out", help="Output file")
    args = parser.parse_args()

    if args.command == "tileset":
        print(f"Processing: {args.input}")
        output_filepath_no_ext = drop_extension(args.input)
        output_filepath = output_filepath_no_ext + ".ts"
        convert_tsx(args.input, output_filepath)
    elif args.command == "tilemap":
        print(f"Processing: {args.input}")
        output_filepath_no_ext = drop_extension(args.input)
        output_filepath = output_filepath_no_ext + ".tm"
        convert_tmx(args.input, output_filepath)
    elif args.command == "sprite":
        print(f"Processing: {args.input}")
        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())