#!/usr/bin/python3
'''
[Tiler] Durden version 1.3 (c) 2024 Carl Svensson, www.datagubbe.se.
Find unique, arbitrarily-sized tiles in a given image and mark and/or
save them in new files.

This program may be distributed and used free of charge.
The program, in whole or in part, may not be modified or
sold without the explicit written consent of the author.
Disclaimer: This program comes without any warranty whatsoever.
            Use it at your own risk.

Arguments:
  <FILENAME>
    The source image for counting, extracting and marking tiles.
    Files are expected to be in non-lossy bitmap formats in RGB, RGBA
    or with a fixed, indexed palette.
  w=<TILE WIDTH>   (Default 16)
    Set tile width to the supplied value.
    Must be a multiple of the source image width.
  h=<TILE HEIGHT>  (Default 16)
    Set tile height to the supplied value.
    Must be a multiple of the source image height.
  wh=<TILE SIZE>   (Default 16)
    Sets tile width and tile height to the supplied value.
    Must be a multiple of the source image height.
  m=<MARK INDEX>   (Default None)
    A palette index in a fixed-palette source image used for
    marking repeating tiles (see below).

Flags:
  mark   Mark repeating tiles in a copy of the source image and save the copy.
         * In a 24-bit RGB image, repeating tiles will be marked with red
           (RGB 255, 0, 0).
         * In an 8-bit indexed image, repeating tiles will be marked with
           the palette index supplied using the "m=" argument.

  alpha  Mark repeating tiles with half opacity alpha channel. Requires
         that the source image is in RGBA format.

  save   Cut unique tiles out and save them. The resulting tiles are
         saved in a directory called "<FILE_NAME>_tiles". If the
         directory doesn't exist, it will be created.

Examples:
  $ durden "my image.png" wh=8 save
      * Find unique 8x8 pixel tiles in the file "my image.png"
      * Save all unique tiles in a directory that will be
        named "my image.png_tiles".

  $ durden "sea_back.png" w=16 h=8 mark m=255
      * Find unique 16x8 pixel tiles in the file "sea_back.png".
      * Mark all repeating tiles with the colour in palette index 255.
      * Save the result in "marked_sea_back.png".

  $ durden "bg_graphics.png" mark alpha
      * Find unique 16x16 pixel tiles in the file "bg_graphics.png".
      * Mark all repeating tiles with half opacity alpha channel.
      * Save the result in "marked_bg_graphics.png".
'''
import os
import sys
from PIL import Image


MARK_INDEXED = 0
MARK_ALPHA = -1
MARK_RED = -2
MARK_NONE = -3


def errout(s):
  sys.stderr.write(f"{s}\n")


def errdie(s):
  errout(s)
  sys.exit(1)


def ensure_numeric(val, val_name):
  if not str(val).isnumeric():
    errdie(f"{val_name} must be numeric, not '{val}'")
  return int(val)


def load_image(file_name):
  if not os.path.exists(file_name):
    errdie(f"Cannot find file '{file_name}'")
  return Image.open(file_name)


def assert_image_props(img, width, height, marker):
  if img.mode not in ["P", "RGB", "RGBA"]:
    errdie("Image is not RGB, RGBA or 8-bit indexed palette.")

  if marker == MARK_ALPHA and img.mode != "RGBA":
    errdie("Mark Alpha: Image lacks alpha channel.")

  if marker >= MARK_INDEXED and img.mode != "P":
    errdie("Mark Index: Image lacks indexed palette.")

  if marker == MARK_RED and img.mode == "P":
    errdie("Mark: Palette index to use for marking not set.")

  if img.width % width:
    errdie(f"Image width {img.width} incompatible with tile size {width}")

  if img.height % height:
    errdie(f"Image height {img.height} incompatible with tile size {height}")


def get_tiles(img, width, height):
  pixels = list(img.getdata())
  pixels = [pixels[(i * img.width) : ((i + 1) * img.width)]
            for i in range(img.height)]

  dupes = {}
  unique = set()
  for row in range(0, img.height, height):
    for col in range(0, img.width, width):
      tile = tuple((tuple(pixels[i][col:col + width])
              for i in range(row, row + height)))
      if tile not in unique:
        unique.add(tile)
      else:
        dupes[(col, row)] = tile

  return unique, dupes


def save_tiles(img, unique_tiles, width, height):
  dir_name = f"{img.filename}_tiles"
  if not os.path.exists(dir_name):
    try:
      os.mkdir(dir_name)
    except os.error as mkd_err:
      errdie(f"Unable to create tiles directory: {mkd_err.strerror}")

  for tilecount, tile in enumerate(unique_tiles):
    tilepic = Image.new(mode=img.mode, size=(width, height))
    if img.palette:
      tilepic.putpalette(img.getpalette())
    for pixy, row in enumerate(tile):
      for pixx, pixel in enumerate(row):
        tilepic.putpixel((pixx, pixy), pixel)
    tilepic.save(f"{img.filename}_tiles/tile_{tilecount}.png")

  return dir_name


def mark_tiles(img, dupes, marker):
  mark_img = img.copy()

  for tile_origin, tile_pixels in dupes.items():
    for pixy, row in enumerate(tile_pixels):
      for pixx, pixel in enumerate(row):
        if img.mode == "P" and marker >= MARK_INDEXED:
          new_pixel = marker
        elif marker == MARK_ALPHA:
          new_pixel = pixel[:3] + (pixel[3]//2,)
        else:
          new_pixel = (255, 0, 0)
        coords = (tile_origin[0] + pixx, tile_origin[1] + pixy)
        mark_img.putpixel(coords, new_pixel)

  save_name = f"marked_{img.filename}"
  try:
    mark_img.save(save_name)
  except IOError as save_err:
    errdie(f"Unable to save marked image: {save_err.strerror}")

  return save_name


def print_help():
  print(__doc__.strip())


def handle_args():
  argv = sys.argv

  help_flags = ["-h", "-H", "--help", "help", "HELP", "?", "/?"]
  for flag in help_flags:
    if flag in argv:
      print_help()
      sys.exit(0)

  filename = None
  width = 16
  height = 16
  index = None
  save = "save" in argv
  mark = "mark" in argv
  alpha = "alpha" in argv

  unified_wh_arg = any(arg.startswith("wh=") for arg in argv)
  w_or_h_args = any("w=" in arg[:2] or "h=" in arg[:2] for arg in argv)
  if unified_wh_arg and w_or_h_args:
    errdie("Cannot specify both wh and h or w.")

  for arg in argv:
    if arg.startswith("w="):
      width = ensure_numeric(arg[2:], "Width")
    elif arg.startswith("h="):
      height = ensure_numeric(arg[2:], "Height")
    elif arg.startswith("wh="):
      width = ensure_numeric(arg[3:], "Width & Height")
      height = ensure_numeric(arg[3:], "Width & Height")
    elif arg.startswith("m="):
      index = ensure_numeric(arg[2:], "Mark index")
    elif arg not in ["alpha", "mark", "save"]:
      script_file = os.path.basename(__file__)
      arg_file = os.path.basename(arg)
      if arg_file != script_file:
        filename = arg

  if not filename:
    errdie("No image name supplied. Use -h for help.")

  return filename, width, height, save, mark, alpha, index


def main():
  filename, width, height, save, mark, alpha, index = handle_args()
  img = load_image(filename)
  marker = (MARK_NONE if not mark else
            MARK_ALPHA if alpha else
            index if index else MARK_RED)

  assert_image_props(img, width, height, marker)

  unique, dupes = get_tiles(img, width, height)
  numtiles = len(unique)
  iname = img.filename
  print(f"Number of unique {width} * {height} tiles in '{iname}': {numtiles}")

  if save:
    savedir = save_tiles(img, unique, width, height)
    print(f"Tiles saved in {savedir}")

  if mark:
    markfile = mark_tiles(img, dupes, marker)
    print(f"Marked image saved as {markfile}")

if __name__ == "__main__":
  main()
