# Shadowrun Huffman Decompressor
# Written by Alchemic
# 2011 Aug 26
# 
# This code uses python-bitstring version 2.2.0:
# http://code.google.com/p/python-bitstring/

import sys
import array
from bitstring import ConstBitStream





# A node class for the Huffman tree.
class HuffmanNode:

    # __init__()
    # Initialization method.
    def __init__(self):
        # Left and right child nodes
        self.left = None
        self.right = None

        # Node data
        self.isLeaf = True
        self.dataLow = 0
        self.dataHigh = 0



    # buildTree()
    # Constructs the Huffman tree.
    def buildTree(self, romStream, index):
        # The pointers are relative to 0xE8000.
        index += 0xE8000

        # This node is now internal.
        self.isLeaf = False

        # Retrieve the data for the left and right child nodes.
        romStream.bytepos = index
        leftData = romStream.read('uintle:16')
        rightData = romStream.read('uintle:16')

        # Left child node.
        self.left = HuffmanNode()
        if bool(leftData & 0x8000) == True:
            self.left.dataHigh = (leftData & 0x7F00) >> 8
            self.left.dataLow = (leftData & 0x007F)
        else:
            self.left.buildTree(romStream, leftData)

        # Right child node.
        self.right = HuffmanNode()
        if bool(rightData & 0x8000) == True:
            self.right.dataHigh = (rightData & 0x7F00) >> 8
            self.right.dataLow = (rightData & 0x007F)
        else:
            self.right.buildTree(romStream, rightData)

        return



    # isEndOfString()
    # Determines whether this node ends the current line.
    def isEndOfString(self):
        if self.isLeaf == False:
            return False
        if self.dataHigh == 0x0A:
            return True
        if self.dataLow == 0x0A:
            return True
        return False



    # __str__()
    # Produces a string representation of the node's data.
    def __str__(self):
        # Start with an empty return string
        returnString = ""

        # Internal nodes have no string representation
        if self.isLeaf == False:
            return returnString

        # First character
        if self.dataHigh == 0x00:
            pass
        elif (self.dataHigh < 0x20) or (self.dataHigh > 0x7E):
            returnString += "{0:02X}".format(self.dataHigh)
        else:
            returnString += "{0:c}".format(self.dataHigh)

        # Second character
        if self.dataLow == 0x00:
            pass
        elif (self.dataLow < 0x20) or (self.dataLow > 0x7E):
            returnString += "{0:02X}".format(self.dataLow)
        else:
            returnString += "{0:c}".format(self.dataLow)

        return returnString







# Check for incorrect usage.
argc = len(sys.argv)
if argc < 3 or argc > 4:
    sys.stdout.write("Usage: ")
    sys.stdout.write("{0:s} ".format(sys.argv[0]))
    sys.stdout.write("<romFile> <offset> [maxLines]\n")
    sys.exit(1)

# Copy the arguments.
romFile = sys.argv[1]
offset = int(sys.argv[2], 16)
maxLines = 1
if argc == 4:
    maxLines = int(sys.argv[3], 10)

# Open the ROM.
romStream = ConstBitStream(filename=romFile)

# Construct the Huffman tree.
rootNode = HuffmanNode()
rootNode.buildTree(romStream, 0)



# Huffman decompression loop.
romStream.bytepos = offset
lineCount = 0

while lineCount < maxLines:
    # Describe the next line.
    sys.stdout.write("{0:4d}    ".format(lineCount))
    sys.stdout.write("{0:5X}    ".format(romStream.bytepos))
    currentNode = rootNode

    # Decompress the next line.
    while currentNode.isEndOfString() == False:
        # Return to the root node.
        currentNode = rootNode
        # Navigate through the internal nodes to reach a leaf.
        while currentNode.isLeaf == False:
            nextBranch = romStream.read('bool')
            if nextBranch == False:
                currentNode = currentNode.left
            else:
                currentNode = currentNode.right
        # Print the leaf node's characters.
        sys.stdout.write(str(currentNode))

    # End the current line.
    romStream.bytealign()
    sys.stdout.write("\n")
    lineCount += 1



# Report the last offset.
sys.stdout.write("        {0:5X}\n".format(romStream.bytepos))

# Exit.
sys.exit(0)