#!/usr/bin/env python
"""Extracts PNG images from a PlayOnline ANG file.

PlayOnline uses 'ANG' files (image/x-playonline-ang) to store multiple images
in a single bundle (e.g., sprites, button states, or animations). These are not
'Animated PNG' (APNG) files (developed about eight years after PlayOnline), but
a custom format that stores frames as a series of PNG images.

To save space, dummy images store common colour palette/transparency chunks,
which are substituted into subsequent images.

NB: The APNG header isn't parsed, so animation or palette association data
    (for files with multiple palettes) isn't read.
"""

from collections import OrderedDict
import sys

import png

PNG_HEADER = '\x89PNG'
PNG_FOOTER = 'IEND\xAE\x42\x60\x82'

if len(sys.argv) != 2:
  print 'Usage: %s ang_file' % sys.argv[0]
  sys.exit(1)
  
ang_data = None
with open(sys.argv[1], 'r') as f:
  ang_data = f.read()

if not ang_data:
  print 'ERROR reading ANG file: %s' % sys.argv[1]
  sys.exit(1)

if ang_data[:4] != '@ANG':
  print 'ERROR: No @ANG header!'
  sys.exit(1)

# Corresponding lists of 'PLTE'/'tRNS' chunks from dummy images
plte_chunks = list()
trns_chunks = list()

base_name = sys.argv[1].rsplit('.', 1)[0]
png_count = 0   # Number of written images

png_start = ang_data.find(PNG_HEADER)
while png_start != -1:
  png_end = ang_data.find(PNG_FOOTER, png_start)
  if png_end == -1:
    print 'ERROR locating end of PNG image from offset %d' % png_start
    sys.exit(1)

  png_data = png.Reader(bytes=ang_data[png_start:png_end + len(PNG_FOOTER)])
  chunks = list(png_data.chunks())
  
  # Concatenate IDAT chunks together
  image_data = ''.join((chunk[1] for chunk in chunks if chunk[0] == 'IDAT'))

  chunks = OrderedDict(chunks)
  chunks['IDAT'] = image_data

  if not image_data:
    # If there's no image data, it's a dummy image
    print 'INFO: Reading palette from image at offset %d' % png_start
    plte_chunks.append(chunks.get('PLTE', None))
    trns_chunks.append(chunks.get('tRNS', None))
  else:
    # Complete image... add in any missing palette/transparency data
    if len(plte_chunks) > 1:
      print 'WARNING: More than one palette available; using first'

    if 'PLTE' in chunks and not chunks['PLTE']:
      if not len(plte_chunks):
        print 'ERROR: No palette data for image at offset %d' % png_start
        sys.exit(1)
      chunks['PLTE'] = plte_chunks[0]

    if 'tRNS' in chunks and not chunks['tRNS']:
      if not len(trns_chunks):
        print 'ERROR: No palette data for image at offset %d' % png_start
        sys.exit(1)
      chunks['tRNS'] = trns_chunks[0]

    file_name = '%s-%d.png' % (base_name, png_count)
    print 'INFO: Writing %s...' % file_name

    png_file = open(file_name, 'w')
    png.write_chunks(png_file, chunks.items())
    png_file.close()

    png_count += 1

  png_start = ang_data.find(PNG_HEADER, png_end)