#!/usr/bin/env python

# s      p     a    r   s  i fy                                            v0.1
#                                                            by itsnotabigtruck
#
# This program is free software. It comes without any warranty, to the extent
# permitted by applicable law. You can redistribute it and/or modify it under
# the terms of the Do What The Fuck You Want To Public License, Version 2, as
# published by Sam Hocevar. See http://sam.zoy.org/wtfpl/COPYING for more
# details.

import itertools
import lzo
import os
import struct
import sys

def append_chunk(data, dest):
    # Emit an LZO-compressed chunk
    m1 = 0xb8c3b410
    m2 = 0
    origLen = len(data)
    if origLen > 0:
        compData = lzo.compress(data)[5:]
        compLen = len(compData)
        if compLen >= origLen:
            compData = data
            compLen = origLen
            comp = 0
        else:
            comp = 1
    else:
        compData = data
        comp = 0
        compLen = origLen
    hdr = struct.pack("<IIIII", m1, m2, comp, compLen, origLen)
    dest.write(hdr)
    dest.write(compData)

def encode_file(name, src, dest):
    # Index the input file for continuous regions of 512-byte blocks
    origPos = src.tell()
    offset = 0
    blocks = []
    start = 0
    count = 0
    totalBlocks = 0
    zData = "\0" * 0x20000
    while True:
        data = src.read(0x20000)
        if data == None:
            break
        dataLen = len(data)
        if dataLen == 0:
            break
        if (dataLen % 512) != 0:
            raise IOError("input file must have multiple of 512 file size")
        if data == zData:
            if count != 0:
                blocks.append((start, count))
                totalBlocks += count
                count = 0
        elif count == 0:
            start = offset
            count = dataLen / 512
        else:
            count += dataLen / 512
        offset += dataLen / 512
    if count != 0:
        blocks.append((start, count))
        totalBlocks += count
    if offset != (start + count):
        blocks.append((offset - 1, 1))
    if len(blocks) > 1016:
        raise IOError("too many sparse regions to represent")
    sys.stderr.write("NAME:           {0}\n".format(name))
    sys.stderr.write("SEGMENTS:       {0}\n".format(len(blocks)))
    sys.stderr.write("BLOCKS READ:    {0}\n".format(offset))
    sys.stderr.write("BLOCKS WRITTEN: {0}\n".format(totalBlocks))

    # Generate a properly formatted header from the block info
    magic = 0xadaea9a1
    hdr = struct.pack("<II56s", magic, len(blocks), name)
    for start, count in blocks:
        hdr += struct.pack("<II", start, count)
    hdr += "\0" * (8 * (1016 - len(blocks)))
    chunk = hdr

    # Emit the header and original data to the output file
    for start, count in blocks:
        src.seek(origPos + (start * 512), 0)
        for _ in itertools.repeat(None, count):
            data = src.read(512)
            if data == None or len(data) != 512:
                raise IOError("failed to transfer input data")
            chunk += data
            if len(chunk) == 0x10000:
                append_chunk(chunk, dest)
                chunk = ""
    if len(chunk) != 0:
        append_chunk(chunk, dest)

    # Append an empty chunk to the end to mark EOF
    append_chunk("", dest)

def main():
    # Validate arguments
    if len(sys.argv) != 3:
        sys.stderr.write("Usage: sparsify.py <input> <output>\n")
        sys.exit(1)
    inPath = sys.argv[1]
    outPath = sys.argv[2]

    # Encode the data
    try:
        inName = os.path.basename(inPath)
        if len(inName) == 0 or len(inName) > 56:
            raise IOError("invalid filename (must be <= 56 characters)")
        with open(inPath, "rb") as inFile, open(outPath, "wb") as outFile:
            encode_file(inName, inFile, outFile)
    except:
        sys.stderr.write("ERROR: " + str(sys.exc_info()[1]) + "\n")
        sys.exit(1)

if __name__ == "__main__":
    main()