#!/usr/bin/python

# system packages
import sys
import os
import argparse
import re

# external packages
import png  # package name is "pypng" on pypi.python.org


class Colors:
    black,magenta,green,orange,blue,white,key = range(7)


def slugify(s):
    """Simplifies ugly strings into something that can be used as an assembler
    label.

    >>> print slugify("[Some] _ Article's Title--")
    SOME_ARTICLES_TITLE

    From https://gist.github.com/dolph/3622892#file-slugify-py
    """

    # "[Some] _ Article's Title--"
    # "[SOME] _ ARTICLE'S TITLE--"
    s = s.upper()

    # "[SOME] _ ARTICLE'S_TITLE--"
    # "[SOME]___ARTICLE'S_TITLE__"
    for c in [' ', '-', '.', '/']:
        s = s.replace(c, '_')

    # "[SOME]___ARTICLE'S_TITLE__"
    # "SOME___ARTICLES_TITLE__"
    s = re.sub('\W', '', s)

    # "SOME___ARTICLES_TITLE__"
    # "SOME   ARTICLES TITLE  "
    s = s.replace('_', ' ')

    # "SOME   ARTICLES TITLE  "
    # "SOME ARTICLES TITLE "
    s = re.sub('\s+', ' ', s)

    # "SOME ARTICLES TITLE "
    # "SOME ARTICLES TITLE"
    s = s.strip()

    # "SOME ARTICLES TITLE"
    # "SOME_ARTICLES_TITLE"
    s = s.replace(' ', '_')

    return s


class Syntax(object):
    def asm(self, text):
        return "\t%s" % text

    def comment(self, text):
        return "\t; %s" % text

    def label(self, text):
        return text

    def byte(self, text):
        return self.asm(".byte %s" % text)

    def word(self, text):
        return self.asm(".word %s" % text)

    def address(self, text):
        return self.asm(".addr %s" % text)

    def origin(self, text):
        return self.asm("*= %s" % text)


class Mac65(Syntax):
    def address(self, text):
        return self.asm(".word %s" % text)


class CC65(Syntax):
    def label(self, text):
        return "%s:" % text


class Listing(object):
    disclaimer = '''
; This file was generated by HiSprite.py, a sprite compiler by Quinn Dunki.
; If you feel the need to modify this file, you are probably doing it wrong.
'''

    def __init__(self, assembler):
        self.assembler = assembler
        self.lines = [self.disclaimer]

    def __str__(self):
        return "\n".join(self.lines) + "\n"

    def out(self, line):
        self.lines.append(line)

    def out_append_last(self, line):
        self.lines[-1] += line

    def label(self, text):
        self.out(self.assembler.label(text))

    def comment(self, text):
        self.out_append_last(self.assembler.comment(text))

    def asm(self, text):
        self.out(self.assembler.asm(text))

    def addr(self, text):
        self.out(self.assembler.address(text))

    def byte(self, text):
        self.out(self.assembler.byte(text))

    def word(self, text):
        self.out(self.assembler.word(text))


class Sprite(Listing):
    def __init__(self, pngfile, assembler, xdraw=False, processor="any"):
        Listing.__init__(self, assembler)

        reader = png.Reader(pngfile)
        try:
            pngdata = reader.asRGB8()
        except:
            raise RuntimeError

        self.xdraw = xdraw
        self.processor = processor
        self.niceName = slugify(os.path.splitext(pngfile)[0])
        self.width = pngdata[0]
        self.height = pngdata[1]
        self.pixelData = list(pngdata[2])
        self.calcStorage()
        self.jumpTable()
        for i in range(self.numShifts):
            self.blitShift(i)

    def calcStorage(self):
        self.byteWidth = self.width/2+1+1  # TODO: Calculate a power of two for this
        self.numShifts = 7

    def jumpTable(self):
        # Prologue
        self.label("%s" % self.niceName)
        self.comment("%d bytes per row" % self.byteWidth)
        self.asm("SAVE_AXY")
        self.asm("ldy PARAM0")
        self.asm("ldx MOD7_2,y")

        if self.processor == "any":
            self.out(".ifpC02")
            self.jump65C02()
            self.out(".else")
            self.jump6502()
            self.out(".endif")
        elif self.processor == "65C02":
            self.jump65C02()
        elif self.processor == "6502":
            self.jump6502()
        else:
            raise RuntimeError("Processor %s not supported" % self.processor)

    def jump65C02(self):
        self.asm("jmp (%s_JMP,x)\n" % (self.niceName))
        offset_suffix = ""
        
        # Bit-shift jump table for 65C02
        self.label("%s_JMP" % (self.niceName))
        for shift in range(self.numShifts):
            self.addr("%s_SHIFT%d" % (self.niceName, shift))

    def jump6502(self):
        # Fast jump table routine; faster and smaller than self-modifying code
        self.asm("lda %s_JMP+1,x" % (self.niceName))
        self.asm("pha")
        self.asm("lda %s_JMP,x" % (self.niceName))
        self.asm("pha")
        self.asm("rts\n")

        # Bit-shift jump table for generic 6502
        self.label("%s_JMP" % (self.niceName))
        for shift in range(self.numShifts):
            self.addr("%s_SHIFT%d-1" % (self.niceName,shift))

    def blitShift(self, shift):
        # Blitting functions
        self.out("\n")
        
        # Track cycle count of the blitter. We start with fixed overhead:
        # SAVE_AXY + RESTORE_AXY + rts +    sprite jump table
        cycleCount = 9 + 12 + 6 +   3 + 4 + 6
    
        self.label("%s_SHIFT%d" % (self.niceName,shift))
        self.asm("ldx PARAM1")
        cycleCount += 3
        rowStartCode,extraCycles = self.rowStartCalculatorCode();
        self.out(rowStartCode)
        cycleCount += extraCycles
        
        spriteChunks = self.layoutSpriteChunk(shift, cycleCount)
        
        for row in range(self.height):
            for chunkIndex in range(len(spriteChunks)):
                self.out(spriteChunks[chunkIndex][row])

    def layoutSpriteChunk(self, shift, cycleCount):

        colorStreams = self.byteStreamsFromPixels(shift, bitsForColor, highBitForColor)
        # print colorStreams
        maskStreams = self.byteStreamsFromPixels(shift, bitsForMask, highBitForMask)
        # print maskStreams
        code = self.generateBlitter(colorStreams, maskStreams, cycleCount)

        return code

    def byteStreamsFromPixels(self, shift, bitDelegate, highBitDelegate):
        byteStreams = ["" for x in range(self.height)]
        byteWidth = self.width/2+1+1

        for row in range(self.height):
            bitStream = ""
            
            # Compute raw bitstream for row from PNG pixels
            for pixelIndex in range(self.width):
                pixel = pixelColor(self.pixelData,row,pixelIndex)
                bitStream += bitDelegate(pixel)
            
            # Shift bit stream as needed
            bitStream = shiftStringRight(bitStream,shift)
            bitStream = bitStream[:byteWidth*8]
            
            # Split bitstream into bytes
            bitPos = 0
            byteSplits = [0 for x in range(byteWidth)]
            
            for byteIndex in range(byteWidth):
                remainingBits = len(bitStream) - bitPos
                    
                bitChunk = ""
                
                if remainingBits < 0:
                    bitChunk = "0000000"
                else:   
                    if remainingBits < 7:
                        bitChunk = bitStream[bitPos:]
                        bitChunk += fillOutByte(7-remainingBits)
                    else:   
                        bitChunk = bitStream[bitPos:bitPos+7]
                
                bitChunk = bitChunk[::-1]
                
                # Determine palette bit from first pixel on each row
                highBit = highBitDelegate(self.pixelData[row][0])
                
                byteSplits[byteIndex] = highBit + bitChunk
                bitPos += 7
                
                byteStreams[row] = byteSplits;

        return byteStreams

    def generateBlitter(self, colorStreams, maskStreams, baseCycleCount):
        byteWidth = len(colorStreams[0])
        spriteChunks = [["" for y in range(self.height)] for x in range(byteWidth)]
        
        cycleCount = baseCycleCount
        optimizationCount = 0

        for row in range(self.height):
            
            byteSplits = colorStreams[row]
            
            # Generate blitting code
            for chunkIndex in range(len(byteSplits)):
                
                # Optimization
                if byteSplits[chunkIndex] != "00000000" and \
                    byteSplits[chunkIndex] != "10000000":
                
                    # Store byte into video memory
                    if self.xdraw:
                        spriteChunks[chunkIndex][row] = \
                        "\tlda (SCRATCH0),y\n" + \
                        "\teor #%%%s\n" % byteSplits[chunkIndex] + \
                        "\tsta (SCRATCH0),y\n";
                        cycleCount += 5 + 2 + 6
                    else:
                        spriteChunks[chunkIndex][row] = \
                        "\tlda #%%%s\n" % byteSplits[chunkIndex] + \
                        "\tsta (SCRATCH0),y\n";
                        cycleCount += 2 + 6
                else:
                    optimizationCount += 1
                
                # Increment indices
                if chunkIndex == len(byteSplits)-1:
                    spriteChunks[chunkIndex][row] += "\n"
                else:   
                    spriteChunks[chunkIndex][row] += "\tiny"
                    cycleCount += 2

            # Finish the row
            if row<self.height-1:
                rowStartCode, extraCycles = self.rowStartCalculatorCode()
                spriteChunks[chunkIndex][row] += "\tinx\n" + rowStartCode;
                cycleCount += 2 + extraCycles
            else:
                spriteChunks[chunkIndex][row] += "\tRESTORE_AXY\n"
                spriteChunks[chunkIndex][row] += "\trts\t;Cycle count: %d, Optimized %d rows." % (cycleCount,optimizationCount) + "\n"
                
        return spriteChunks

    def rowStartCalculatorCode(self):
        return \
        "\tlda HGRROWS_H1,x\n" + \
        "\tsta SCRATCH1\n" + \
        "\tlda HGRROWS_L,x\n" + \
        "\tsta SCRATCH0\n" + \
        "\tldy PARAM0\n" + \
        "\tlda DIV7_2,y\n" + \
        "\ttay\n", 4 + 3 + 4 + 3 + 3 + 4 + 2;


def fillOutByte(numBits):
    filler = ""
    for bit in range(numBits):
        filler += "0"
    
    return filler


def shiftStringRight(string,shift):
    if shift==0:
        return string
    
    shift *=2   
    result = ""
    
    for i in range(shift):
        result += "0"
        
    result += string
    return result
                

def bitsForColor(pixel):
    if pixel == Colors.black:
        return "00"
    else:
        if pixel == Colors.white:
            return "11"
        else:
            if pixel == Colors.green or pixel == Colors.orange:
                return "01"

    # blue or magenta
    return "10"


def bitsForMask(pixel):
    if pixel == Colors.black:
        return "00"

    return "11"


def highBitForColor(pixel):
    # Note that we prefer high-bit white because blue fringe is less noticeable than magenta.
    highBit = "0"
    if pixel == Colors.orange or pixel == Colors.blue or pixel == Colors.white:
        highBit = "1"

    return highBit


def highBitForMask(pixel):
    return "1"


def pixelColor(pixelData,row,col):
    r = pixelData[row][col*3]
    g = pixelData[row][col*3+1]
    b = pixelData[row][col*3+2]
    color = Colors.black
    
    if r==255 and g==0 and b==255:
        color = Colors.magenta
    else:
        if r==0 and g==255 and b==0:
            color = Colors.green
        else:
            if r==0 and g==0 and b==255:
                color = Colors.blue
            else:
                if r==255 and g>0 and b==0:
                    color = Colors.orange
                else:
                    if r==255 and g==255 and b==255:
                        color = Colors.white
                    else:
                        if r==g and r==b and r!=0 and r!=255:   # Any gray is chroma key
                            color = Colors.key
    return color


class HorizontalLookup(Listing):
    def __init__(self, assembler):
        Listing.__init__(self, assembler)
        self.generate_hgr()
        self.generate_tables()

    def generate_hgr(self):
        offsets = []
        for y in range(192):
            # From Apple Graphics and Arcade Game Design
            a = y // 64
            d = y - (64 * a)
            b = d // 8
            c = d - 8 * b
            offsets.append((1024 * c) + (128 * b) + (40 * a))

        self.label("HGRROWS_H1")
        for y in range(192):
            addr = 0x2000 + offsets[y]
            self.byte("$%02x" % (addr // 256))

        self.label("HGRROWS_H2")
        for y in range(192):
            addr = 0x4000 + offsets[y]
            self.byte("$%02x" % (addr // 256))

        self.label("HGRROWS_L")
        for y in range(192):
            addr = offsets[y]
            self.byte("$%02x" % (addr & 0xff))

    def generate_tables(self):
        self.label("DIV7_2")
        for pixel in range(140):
            self.byte("$%02x" % ((pixel / 7)*2))

        self.out("\n")
        self.label("MOD7_2")
        for pixel in range(140):
            self.byte("$%02x" % ((pixel % 7)*2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Sprite compiler for 65C02/6502 to generate assembly code to render all shifts of the given sprite, optionally with exclusive-or drawing (if background will be non-black). Generated code has conditional compilation directives for the CC65 assembler to allow the same file to be compiled for either architecture.")
    parser.add_argument("-v", "--verbose", default=0, action="count")
    parser.add_argument("-t", "--tables", action="store_true", default=False, help="output only lookup tables for horizontal sprite shifts (division and modulus 7)")
    parser.add_argument("-x", "--xdraw", action="store_true", default=False, help="use XOR for sprite drawing")
    parser.add_argument("-s", "--syntax", default="cc65", nargs=1, choices=["cc65","mac65"], help="Assembler syntax (default: %(default)s)")
    parser.add_argument("-p", "--processor", default="any", nargs=1, choices=["any","6502", "65C02"], help="Processor type (default: %(default)s)")
    parser.add_argument("files", metavar="IMAGE", nargs="*", help="a PNG image [or a list of them]. PNG files must not have an alpha channel!")
    options, extra_args = parser.parse_known_args()

    if options.syntax[0].lower() == "cc65":
        syntax = CC65()
    elif options.syntax[0].lower() == "mac65":
        syntax = Mac65()
    else:
        print("Unknown assembler %s" % options.syntax)
        parser.print_help()

    if options.tables:
        print HorizontalLookup(syntax)
        exit(0)

    for pngfile in options.files:
        try:
            print Sprite(pngfile, syntax, options.xdraw, options.processor[0])
        except RuntimeError, e:
            print e
            parser.print_help()