# Shadowrun Decompressor
# Written by Alchemic
# 2011 Aug 25
# 
# 
# 
# A detailed description of the compression format:
# 
#   - Compressed data consists of three parts: the header, the data 
#     section, and the control section.
# 
#   - The header consists of two 16-bit integers. The first integer 
#     contains the length of the data once decompressed. The second 
#     indicates where the control section starts: add the value of 
#     the second integer to the address where it's located and you 
#     have the address where the control section begins.
# 
#   - For example, here's the header for the Oldtown Magic Shop's 
#     compressed drawing data (0x5F163):
# 
#        0x5F163  [8D 00] --> 0x008D
#        0x5F165  [6B 00] --> 0x006B
# 
#        There are 0x008D (141) bytes in the decompressed output.
#        The control section begins at 0x5F165 + 0x006B = 0x5F1D0.
# 
#   - The data section occupies all of the space after the header 
#     and before the control section. (So if you take the second 
#     integer from the header and subtract 2, you get the size in 
#     bytes of the data section.) It is read strictly sequentially
#     by literal commands, which we will see shortly.
# 
#   - Last is the control section: a stream of bits that breaks
#     down into two commands, "literal" and "pastcopy". The bits 
#     are read from one byte at a time, most significant to least 
#     (0x80, 0x40, 0x20 ... 0x01). The first command is always 
#     literal.
# 
#   - The literal command copies a specified number of bytes 
#     from the data section into the decompressed output.
#     To read a literal:
#      - Set the initial amount to 1.
#      - Read one bit.
#      - If the bit is 0, the literal continues. Shift the 
#        amount left once, add the next bit to the amount, 
#        and return to the previous step.
#      - If the bit is 1, the literal is done. Copy the 
#        specified amount of bytes from the data section 
#        into the output.
#     Literals are always followed by a pastcopy.
# 
#   - Some sample literals:
# 
#                  1 --> Copy  1 byte  (binary      1)
#                001 --> Copy  2 bytes (binary     10)
#                011 --> Copy  3 bytes (binary     11)
#              00001 --> Copy  4 bytes (binary    100)
#              00011 --> Copy  5 bytes (binary    101)
#              01001 --> Copy  6 bytes (binary    110)
#              01011 --> Copy  7 bytes (binary    111)
#            0000001 --> Copy  8 bytes (binary   1000)
#        01000100001 --> Copy 52 bytes (binary 110100)
# 
#   - The pastcopy command consists of three parts.
#     The first is the source: an uninterrupted sequence of bits
#     of length N, where N = log2(number of bytes decompressed),
#     indicating where to copy from.
#     The second is the amount, which is encoded like the amount 
#     in the literal case, except that you add 2 when done.
#     The third is the successor, which is a single bit. If it's
#     0, the next command is a literal, and if it's 1, the next 
#     command is another pastcopy.
# 
# 
# 
# This code uses python-bitstring version 2.2.0:
# http://code.google.com/p/python-bitstring/

import sys
import array
from bitstring import ConstBitStream



# 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> <startOffset> [outFile]\n")
    sys.exit(1)

# Copy the arguments.
romFile = sys.argv[1]
startOffset = int(sys.argv[2], 16)
outFile = None
if argc == 4:
    outFile = sys.argv[3]



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

# Read the header.
decompSize = romStream.read('uintle:16')
decomp = array.array('B', [0x00] * decompSize)
decompPos = 0
dataSize = romStream.read('uintle:16') - 2
data = array.array('B', [0x00] * dataSize)
dataPos = 0

# Read the data section.
for i in range(dataSize):
    data[i] = romStream.read('uint:8')

# Create the log2 table, because rounding errors are not fun.
LOG2 = [
  0, 1, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4,
  5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 
  6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 
  6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 
  7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 
  7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 
  7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 
  7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 
  8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8
]



# Main decompression loop.
nextCommand = False
while decompPos < decompSize:

    if nextCommand == False:
        # 0: Literal case.

        # Read the number of bytes to copy.
        copyAmount = 1
        stopBit = romStream.read('bool')
        while stopBit == False:
            copyAmount <<= 1
            dataBit = int(romStream.read('bool'))
            copyAmount += dataBit
            stopBit = romStream.read('bool')

        # Truncate the copy if it would exceed decompSize.
        if (decompPos + copyAmount) >= decompSize:
            copyAmount = decompSize - decompPos

        # Copy the bytes.
        for i in range(copyAmount):
            decomp[decompPos] = data[dataPos]
            decompPos += 1
            dataPos += 1

        # Literals are always followed by a pastcopy.
        nextCommand = True
        continue

    else:
        # 1: Pastcopy case.

        # Calculate the length of the source value.
        copySourceLength = 0
        if decompPos >= 0x100:
            copySourceLength = 8 + LOG2[decompPos >> 8]
        else:
            copySourceLength = LOG2[decompPos]

        # Read the source.
        copySource = romStream.read('uint:{0:d}'.format(copySourceLength))

        # Read the amount.
        copyAmount = 1
        stopBit = romStream.read('bool')
        while stopBit == False:
            copyAmount <<= 1
            dataBit = int(romStream.read('bool'))
            copyAmount += dataBit
            stopBit = romStream.read('bool')
        copyAmount += 2

        # Truncate the copy if it would exceed decompSize.
        if (decompPos + copyAmount) >= decompSize:
            copyAmount = decompSize - decompPos

        # Copy the bytes.
        for i in range(copyAmount):
            decomp[decompPos] = decomp[copySource]
            decompPos += 1
            copySource += 1

        # If we're not done, read the successor bit.
        if decompPos != decompSize:
            nextCommand = romStream.read('bool')

        # Return to the top of the loop.
        continue



# Report the last offset.
romStream.bytealign()
lastOffset = romStream.bytepos - 1
sys.stdout.write("Last offset read, inclusive: {0:X}\n".format(lastOffset))

# Write the decompressed output, if appropriate.
if outFile is not None:
    outStream = open(outFile, "wb")
    decomp.tofile(outStream)
    outStream.close()

# Exit.
sys.exit(0)