#!/usr/bin/env python3

import glob
import json
import os
import pprint  # pylint: disable=unused-import
import re
import sys

import pygame

# The guide in tiles/guide.png is quite helpful.

# TODO: sort the stuff in free/ and noncommercial/
# TODO: in future/100/characters there are some doors, conveyor belts, screens

SPRITE_FILES = [
    'animals/sheets/*.png',
    'beasttribes/100/*.png',
    'characters/sheets/*.png',
    'christmas/1x/gnome*.png',
    'christmas/1x/reindeer*.png',
    'christmas/1x/rudolph*.png',
    'christmas/1x/xmas*.png',
    'dwarvesvselves/regularsize/*.png',
    'future/100/characters/cars.png',
    'future/100/characters/future*.png',
    'future/100/characters/military*.png',
    'future/100/characters/modern*.png',
    'halloween/ghost1.png',
    'halloween/horseman/*1.png',
    'halloween/reaper/*1.png',
    'halloween/witch/1x/*.png',
    'lichcrusades/100/*.png',
    'monsters/1x/*.png',
    'mythicalbosses/100/*.png',
    'mythicalbosses/dinosaurs/*.png',
    'npcanimations/rpgmaker/1/*.png',
    'ship/100/char/airship*.png',
    'ship/100/char/boat*.png',
    'ship/100/char/pirates_100.png',
    'ship/100/char/ship*.png',
]

SPRITE_SIDEVIEW_FILES = [
    'beasttribes/100/sv_battler/*.png',
    'future/100/svbattler/*.png',
    # TODO: these need to get scaled down 2x before they can be used.
    'sv_battle/RMMV/sv_actors/*.png',
]

TILESET_FILES = [
    'ashlands/ashlands_tileset.png',
    'atlantis/tf_atlantis_tiles.png',
    'beach/beach_tileset.png',
    'christmas/1x/addon_igloo_1.png',
    'christmas/1x/christmas*.png',
    'cloud/cloud_tileset.png',
    'darkdimension/tf_darkdimension_sheet.png',
    'farmandfort/ff_master_tile_sheet.png',
    'future/100/tilesets/*.png',
    'gianttree/tf_gianttree_tiles.png',
    'halloween/tiles/*1.png',
    'jungle/tf_jungle_tileset.png',
    'patron/train_sheet_1.png',
    'ruindungeons/ruindungeons_sheet_full.png',
    'ship/ship_big_tileset.png',
    'tiles/TILESETS/*.png',
    'winter/tiles/*.png',
]

ANIMATION_FILES = [
    'patron/fireworks*_1.png',
    'pixelanimations/animationsheets/*.png',
    'ship/100/char/!$ship_wave*.png',
    'sv_battle/RMMV/system/States.png',
    'tiles/TILESETS/animated/*.png',
]

ICON_FILES = [
    'farmandfort/IconSet/tf_icon_16.png',
    'halloween/hallowicons_1.png',
]

BACKGROUND_FILES = [
    'cloud/bg_*.png',
    'future/100/other/spacebg.png',
    'ship/100/parallax/*.png'
]


def unglob(list_of_globs):
  result = []
  for file in list_of_globs:
    globbed_files = glob.glob(file)
    assert globbed_files, 'glob for %s should be non-empty' % file
    result.extend(globbed_files)
  result.sort()
  return result


def input_wh(prompt):
  while True:
    geometry = input(prompt).strip()
    try:
      cols, rows = [int(x) for x in geometry.split(' ')]
      return cols, rows
    except ValueError:
      pass


# Returns True or False.
def input_ok(prompt):
  while True:
    ok = input(prompt).strip()
    if ok.startswith('y'):
      return True
    if ok.startswith('n'):
      return False


def draw_checkerboard(size):
  surface = pygame.display.get_surface()
  surface.fill((224, 224, 224))
  for i in range(surface.get_width() // size + 1):
    for j in range(surface.get_height() // size + 1):
      if (i + j) % 2 == 0:
        continue
      rect = pygame.Rect(i * size, j * size, size, size)
      surface.fill((192, 192, 192), rect)


def show_splits(image_width, image_height, cols, rows):
  surface = pygame.display.get_surface()
  split_width = image_width / cols
  split_height = image_height / rows
  for i in range(cols):
    for j in range(rows):
      rect = pygame.Rect(
          i * split_width, j * split_height, split_width + 1, split_height + 1)
      pygame.draw.rect(surface, (255, 0, 255), rect, 1)
  pygame.display.flip()


def render_text(text, pos, color):
  surface = pygame.display.get_surface()
  font = pygame.font.SysFont('notomono', 16)
  image = font.render(text, True, color)
  surface.blit(image, pos)
  pygame.display.flip()


def render_sprite(metadata):
  line_color = (255, 0, 255)

  surface = pygame.display.get_surface()
  draw_checkerboard(8)
  image = pygame.image.load(metadata['filename'])
  surface.blit(image, (0, 0))

  if metadata.get('chunks'):
    for chunk in metadata['chunks']:
      rect = pygame.Rect(
          chunk['x'], chunk['y'], chunk['width'] + 1, chunk['height'] + 1)
      pygame.draw.rect(surface, line_color, rect, 1)
      label_pos = (chunk['x'] + 4, chunk['y'])
      render_text(str(chunk['index']), label_pos, line_color)
      caption_pos = (4, 4 + metadata['image_height'] + chunk['index'] * 20)
      caption = '%d: %s' % (chunk['index'], chunk.get('name', ''))
      render_text(caption, caption_pos, (0, 0, 0))

  pygame.display.flip()


def set_sprite_chunk_size(metadata):
  cols, rows = input_wh('how many columns & rows of sprites? ')
  metadata['chunk_columns'] = cols
  metadata['chunk_rows'] = rows
  metadata['chunk_width'] = metadata['image_width'] // cols
  metadata['chunk_height'] = metadata['image_height'] // rows
  metadata['chunks'] = []
  for i in range(cols * rows):
    x = i % cols
    y = i // cols
    chunk_md = {
        'index': i,
        'x': x * metadata['chunk_width'],
        'y': y * metadata['chunk_height'],
        'width': metadata['chunk_width'],
        'height': metadata['chunk_height']
    }
    metadata['chunks'].append(chunk_md)
  render_sprite(metadata)


def edit_sprite_chunk_metadata(chunk):
  while True:
    name = input('name for chunk #%d: ' % chunk['index']).strip()
    if re.fullmatch(r'\w+', name):
      chunk['name'] = name
      return


def edit_sprite_metadata(filename, metadata=None):
  if metadata is None:
    image = pygame.image.load(filename)
    metadata = {
        'filename': filename,
        'image_width': image.get_width(),
        'image_height': image.get_height(),
    }

  print('\nprocessing %s (%dx%d)' % (
      filename, metadata['image_width'], metadata['image_height']))

  render_sprite(metadata)

  if not metadata.get('chunk_width'):
    set_sprite_chunk_size(metadata)

  while True:
    render_sprite(metadata)
    prompt = 'edit (c)hunk sizes, type a chunk #, (n)ext, or (q)uit: '
    choice = input(prompt).strip()
    if choice == 'n':
      return metadata, False
    elif choice == 'q':
      return metadata, True
    elif choice == 'c':
      set_sprite_chunk_size(metadata)
    elif re.fullmatch(r'\d+', choice):
      chunk_num = int(choice)
      if 0 <= chunk_num < len(metadata['chunks']):
        edit_sprite_chunk_metadata(metadata['chunks'][chunk_num])
      else:
        print('invalid chunk #')
    else:
      print('invalid choice')


def annotate_sprites(all_metadata, sprite_files):
  pygame.init()
  pygame.display.set_mode((1200, 900), pygame.RESIZABLE)

  for filename in sprite_files:
    sprite_metadata, should_quit = edit_sprite_metadata(
        filename, all_metadata.get(filename))
    all_metadata[filename] = sprite_metadata
    with open('sprites.json', 'w') as f:
      json.dump(all_metadata, f, sort_keys=True, indent=2)
    if should_quit:
      return


def get_named_sprites(metadata):
  result = {}
  for filename in metadata:
    sprite = metadata[filename]
    for chunk in sprite.get('chunks', []):
      name = chunk.get('name', '')
      if not name:
        continue
      if name in result:
        print('warning: duplicated sprite name ', name)
        sys.exit(1)
      result[name] = chunk
  return result


def check_sprites(metadata):
  named_sprites = get_named_sprites(metadata)
  print('# named sprites:', len(named_sprites))
  with open('sprites.json', 'w') as f:
    json.dump(metadata, f, sort_keys=True, indent=2)


def stitch_sprites(metadata, filename_base, sprite_names=None):
  sprites = get_named_sprites(metadata)
  if sprite_names:
    sprites = dict([(x[0], x[1])
                    for x in sprites.items() if x[0] in sprite_names])
  max_height = 0
  total_width = 0
  for sprite_name, sprite in sprites.items():
    total_width += sprite['width']
    max_height = max(max_height, sprite['height'])
  print('\n# named sprites:', len(sprites))
  print('result will be %dx%d' % (total_width, max_height))

  output = pygame.surface.Surface(
      (total_width, max_height), flags=pygame.SRCALPHA)
  output_json = {}
  xpos = 0
  for sprite_name, sprite in sprites.items():
    sprite_image = pygame.image.load(sprite['filename'])
    area = pygame.Rect(
        sprite['x'], sprite['y'], sprite['width'], sprite['height'])
    output_json[sprite_name] = {
        'name': sprite['name'],
        'x': xpos,
        'y': 0,
        'width': sprite['width'],
        'height': sprite['height']
    }
    output.blit(sprite_image, (xpos, 0), area)
    xpos += sprite['width']
  image_filename = os.path.expanduser(filename_base) + '.png'
  print('saving image to', image_filename)
  pygame.image.save(output, image_filename)
  json_filename = os.path.expanduser(filename_base) + '.js'
  print('saving json to', json_filename)
  with open(json_filename, 'w') as json_file:
    json_file.write('const spritesheet_json = ')
    json.dump(output_json, json_file, sort_keys=True, indent=2)
    json_file.write(';')


def main(args):
  snej_root = os.environ.get('SNEJ_ROOT', os.path.join('~', 'snej'))
  time_fantasy_path = os.path.join(snej_root, 'assets', 'time_fantasy')
  os.chdir(os.path.expanduser(time_fantasy_path))

  sprite_files = unglob(SPRITE_FILES)
  tileset_files = unglob(TILESET_FILES)
  animation_files = unglob(ANIMATION_FILES)
  icon_files = unglob(ICON_FILES)
  background_files = unglob(BACKGROUND_FILES)
  print('\nsprites: %d tilesets: %d animations: %d icons: %d backgrounds: %d' %
        (len(sprite_files), len(tileset_files), len(animation_files),
         len(icon_files), len(background_files)))

  if len(args) < 1:
    return

  command = args[0]

  with open('sprites.json') as f:
    metadata = json.load(f)

  if command == 'annotate-sprites':
    annotate_sprites(metadata, sprite_files)
  elif command == 'check-sprites':
    check_sprites(metadata)
  elif command == 'stitch-sprites':
    if len(args) < 2:
      print('need FILENAME_BASE')
      return
    filename_base = args[1]
    sprite_names = args[2:]
    stitch_sprites(metadata, filename_base, sprite_names)
  else:
    print('unrecognized command "%s"' % command)


if __name__ == '__main__':
  main(sys.argv[1:])