#!/usr/bin/python3

import array
import configparser
import gif
import itertools
import sys

def get_color (color_table, index):
    if index < len (color_table):
        (red, green, blue) = color_table[index]
        return (red, green, blue, 255)
    else:
        return (0, 0, 0, 255)

def set_pixel (reader, pixels, x, y, color):
    offset = (y * reader.width + x) * 4
    (red, green, blue, alpha) = color
    pixels[offset + 0] = red
    pixels[offset + 1] = green
    pixels[offset + 2] = blue
    pixels[offset + 3] = alpha

def get_pixel (reader, pixels, x, y):
    offset = (y * reader.width + x) * 4
    return (pixels[offset + 0], pixels[offset + 1], pixels[offset + 2], pixels[offset + 3])

def render_block (reader, pixels, block, transparent_color):
    i = 0
    block_pixels = block.get_pixels ()
    if len (block.color_table) > 0:
        color_table = block.color_table
    else:
        color_table = reader.color_table
    if block.interlace:
        top = block.top
        bottom = block.top + block.height
        row_iter = itertools.chain (range (top, bottom, 8), range (top + 4, bottom, 8), range (top + 2, bottom, 4), range (top + 1, bottom, 2))
    else:
        row_iter = range (block.top, block.top + block.height)
    for y in row_iter:
        for x in range (block.left, block.left + block.width):
            if i >= len (block_pixels):
                return
            index = block_pixels[i]
            i += 1
            if 0 <= x < reader.width and 0 <= y < reader.height:
                if index == transparent_color:
                    color = (0, 0, 0, 0)
                else:
                    color = get_color (color_table, index)
                set_pixel (reader, pixels, x, y, color)

def dispose_block (reader, pixels, block, method, previous_pixels):
    if method == gif.DisposalMethod.KEEP:
        pass
    elif method == gif.DisposalMethod.RESTORE_BACKGROUND:
        color = (0, 0, 0, 0)
        for y in range (block.top, block.top + block.height):
            for x in range (block.left, block.left + block.width):
                if 0 <= x < reader.width and 0 <= y < reader.height:
                    set_pixel (reader, pixels, x, y, color)
    elif method == gif.DisposalMethod.RESTORE_PREVIOUS:
        for y in range (block.top, block.top + block.height):
            for x in range (block.left, block.left + block.width):
                if 0 <= x < reader.width and 0 <= y < reader.height:
                    color = get_pixel (reader, previous_pixels, x, y)
                    set_pixel (reader, pixels, x, y, color)
        pass # FIXME

def render (reader):
    # Fill background
    pixels = array.array ('B')
    color = (0, 0, 0, 0)
    for y in range (reader.height):
        for x in range (reader.width):
            (red, green, blue, alpha) = color
            pixels.append (red)
            pixels.append (green)
            pixels.append (blue)
            pixels.append (alpha)

    # Write images onto background
    disposal_method = gif.DisposalMethod.NONE
    transparent_color = None
    last_block = None
    previous_pixels = pixels[:]
    for block in reader.blocks:
        if isinstance (block, gif.GraphicControlExtension):
            disposal_method = block.disposal_method
            if block.has_transparent:
                transparent_color = block.transparent_color
        elif isinstance (block, gif.Image):
            if last_block is not None:
                (m, b) = last_block
                dispose_block (reader, pixels, b, m, previous_pixels)
            render_block (reader, pixels, block, transparent_color)
            if disposal_method != gif.DisposalMethod.RESTORE_PREVIOUS:
                previous_pixels = pixels[:]
            last_block = (disposal_method, block)
            disposal_method = gif.DisposalMethod.NONE
            transparent_color = None

    return (reader.width, reader.height, pixels)

def compare_to_reference_frame (reader, pixels, filename):
    r_pixels = open (filename, 'rb').read ()
    if len (pixels) != len (r_pixels):
        print ('  Pixel length mismatch! Got %s, expected %s' % (len (r_pixels), len (pixels)))
        return False
    for y in range (reader.height):
        for x in range (reader.width):
            color = get_pixel (reader, pixels, x, y)
            r_color = get_pixel (reader, r_pixels, x, y)
            if color[3] == r_color[3] == 0:
                pass
            elif color != r_color:
                print ('  Pixel mismatch at %d,%d! Got %s, expected %s' % (x, y, color, r_color))
                return False
    return True;

def run_test (name):
    config = configparser.ConfigParser ()
    config_filename = 'test-suite/%s.conf' % name
    config.read (config_filename);
    if not config.has_section ('config'):
        print ('  Test config %s does not exist / invalid' % config_filename)
        return False

    c = config['config']
    input_filename = c['input']
    frames = []
    frame_names = c['frames'].split (',')
    for frame_name in frame_names:
        if frame_name != '':
            frames.append (config[frame_name])

    print ('  Loading image %s' % input_filename)
    reader = gif.Reader ()
    data = open ('test-suite/%s' % input_filename, 'rb').read ()
    reader.feed (data)

    expected_version = bytes (c['version'], 'utf-8')
    if reader.version != expected_version:
        print ('  Version mismatch!')
        print ('  Got     : %s' % reader.version)
        print ('  Expected: %s' % expected_version)

    expected_width = int (c['width'])
    expected_height = int (c['height'])
    if (expected_width, expected_height) != (reader.width, reader.height):
        print ('  Size mismatch! Got %dx%d, expected %dx%d' % (reader.width, reader.height, expected_width, expected_height))
        return False

    def parse_string (value):
        if value is None:
            return None
        result = ''
        quote = value[0]
        if not quote in ("'", '"'):
            return value
        value = value[1:]
        while len (value) > 1:
            if value.startswith ('\\x'):
                result += chr (int (value[2:4]))
                value = value[4:]
            else:
                result += value[0]
                value = value[1:]
        assert (value[0] == quote)
        return result

    expected_background = parse_string (c.get ('background'))
    if reader.background_color < len (reader.color_table):
        (red, green, blue) = reader.color_table[reader.background_color]
        background = '#%02x%02x%02x' % (red, green, blue)
    else:
        background = None
    if background != expected_background:
        print ('  Background mismatch!')
        print ('  Got     : %s' % repr (background))
        print ('  Expected: %s' % repr (expected_background))
        return False

    expected_comment = parse_string (c.get ('comment'))
    comment = None
    for block in reader.blocks:
        if isinstance (block, gif.CommentExtension):
            comment = block.get_comment ()
    if comment != expected_comment:
        print ('  Comment mismatch!')
        print ('  Got      (%d): %s' % (len (comment), repr (comment)))
        print ('  Expected (%d): %s' % (len (expected_comment), repr (expected_comment)))
        return False

    expected_loop_count = c['loop-count']
    if expected_loop_count != 'infinite':
        expected_loop_count = int (expected_loop_count)
    loop_count = 0
    for block in reader.blocks:
        if isinstance (block, gif.NetscapeExtension):
            if block.loop_count is not None:
                loop_count = block.loop_count
                if loop_count == 0:
                    loop_count = 'infinite'
        if isinstance (block, gif.AnimationExtension):
            if block.loop_count is not None:
                loop_count = block.loop_count
                if loop_count == 0:
                    loop_count = 'infinite'
    if loop_count != expected_loop_count:
        print ('  Loop count mismatch!')
        print ('  Got     : %s' % str (loop_count))
        print ('  Expected: %s' % str (expected_loop_count))
        return False

    value = c.get ('buffer-size')
    if value is None:
        expected_buffer_size = None
    else:
        expected_buffer_size = int (value)
    buffer_size = None
    for block in reader.blocks:
        if isinstance (block, gif.NetscapeExtension):
            if block.buffer_size is not None:
                buffer_size = block.buffer_size
        if isinstance (block, gif.AnimationExtension):
            if block.buffer_size is not None:
                buffer_size = block.buffer_size
    if buffer_size != expected_buffer_size:
        print ('  Buffer size mismatch!')
        print ('  Got     : %d' % buffer_size)
        print ('  Expected: %d' % expected_buffer_size)
        return False

    xmp_filename = c.get ('xmp-data')
    expected_data = ''
    if xmp_filename is not None:
        expected_data = open ('test-suite/%s' % xmp_filename.strip ()).read ()
    xmp_data = ''
    for block in reader.blocks:
        if isinstance (block, gif.XMPDataExtension):
            xmp_data = block.get_metadata ()
    if xmp_data != expected_data:
        print ('  XMP Data mismatch!')
        print ('  Got:')
        print (xmp_data)
        print ('  Expected:')
        print (expected_data)
        return False

    icc_filename = c.get ('color-profile')
    expected_data = b''
    if icc_filename is not None:
        expected_data = open ('test-suite/%s' % icc_filename.strip (), 'rb').read ()
    icc_data = b''
    for block in reader.blocks:
        if isinstance (block, gif.ICCColorProfileExtension):
            icc_data = block.get_icc_profile ()
    if icc_data != expected_data:
        print ('  ICC Color Profile mismatch!')
        print ('  Got     : %s' % repr (icc_data))
        print ('  Expected: %s' % repr (expected_data))
        return False

    # Skip test that uses too much memory
    if name == 'max-size':
        return True

    if len (frames) == 0:
        return True

    (width, height, pixels) = render (reader)

    frame = frames[-1]
    reference_filename = 'test-suite/%s' % frame['pixels']

    print ('  Comparing to %s' % reference_filename)
    return compare_to_reference_frame (reader, pixels, reference_filename)

if len (sys.argv) > 1:
    tests = sys.argv[1:]
else:
    lines = open ('test-suite/TESTS').readlines ()
    tests = []
    for line in lines:
        name = line.strip ()
        if name != '':
            tests.append (name)

successes = []
failures = []
for name in tests:
    print ('Running %s' % name)

    # Skip 87a animation for now - we don't have the animation heuristic
    if name == 'gif87a-animation':
        print ('  SKIP')
    elif run_test (name):
        print ('  PASS')
        successes.append (name)
    else:
        print ('  FAIL')
        failures.append (name)
print ('------------------')
print ('%d/%d tests passed' % (len (successes), len (successes) + len (failures)))
if len (failures) > 0:
    print ('Failures: %s' % ', '.join (failures))
    print ('FAIL')
else:
    print ('PASS')
