#!/usr/bin/env python
"""
An un-optimized but functional script for swapping the audio-track
byte order in cdrdao .BIN files so that a .CUE file generated by toc2cue can
be mounted by DOSBox or CDEmu without the audio tracks being noise.

Now with more sanity checks.

TODO:
- Figure out how I broke this while rewriting it.
  (It's not yet returning output equal to the old version)
- Confirm IO-boundedness for this new design.
- Optimize and, if necessary, rewrite in C
"""

__appname__ = "AudioSwab"
__author__  = "Stephan Sokolow (deitarion/SSokolow)"
__version__ = "0.2alpha0 (broken)"
__license__ = "GNU GPL 2.0 or newer"

CHUNK_SIZE = 4096   #TODO: Benchmark and optimize this value.

import os, shlex, sys

def time_to_frames(timecode):
    """Convert a CUE timecode to a count of BIN frames."""
    mins, secs, frames = [int(x) for x in timecode.split(':')]
    return 75 * (mins * 60 + secs) + frames

def get_extents(cue_path):
    """
    Given a file-like object pointing to a CUE file, return a tuple containing
    the path to the BIN file it describes and a list of (start, end) tuples
    describing the extents of the audio tracks within.
    """
    binpath, offset, track_type, tracks = None, None, None, []

    # Read the extents of the audio tracks from the file
    for line in file(cue_path, 'rU'):
        line = line.strip()
        if line.startswith('FILE'):
            if binpath:
                # I've never found a CUE file which references multiple BIN files.
                raise Exception("Given cuesheet contains references to multiple files.")
            else:
                binfile, mode = shlex.split(line)[1:3]
                binpath = os.path.join(os.path.dirname(cue_path), binfile)
                if mode != 'BINARY':
                    # This tool is only for byte-swapping .bin audio tracks.
                    raise Exception("Given cuesheet does not reference a BINARY source file.")
        elif line.startswith('TRACK'):
            previous_type, track_type = track_type, line.split()[2]
        elif line.startswith('INDEX'):
            previous_offset, offset = offset, time_to_frames(line.split()[2]) * 2352
            if previous_type == 'AUDIO':
                tracks.append((previous_offset, offset))
                # Build a list of start-stop pairs for audio tracks

    # Make sure we don't omit the final track
    if track_type == 'AUDIO':
        tracks.append((offset, os.stat(binpath).st_size))

    # Simplify the problem by combining adjacent extents
    # (In all sane CUEs, this will produce a single giant extent)
    temp = []
    while tracks:
        start, stop = tracks.pop(0)
        if tracks and stop == tracks[0][0]:
            stop = tracks[0][1]
            tracks.pop(0)
    temp.append((start, stop))
    tracks = temp

    return binpath, tracks

def swap_copy(inpath, outpath, extents):
    """
    Given a pair of file paths and a list of audio-track extents, copy the data,
    byte-swapping the audio tracks. (Assumes CD BIN files)
    """
    try:
        infile, outfile = file(inpath, 'rb'), file(outpath, 'wb')

        # Sort the extents (required) and sanity-check them.
        extents.sort(key=lambda x: x[0])
        for pos in range(0, len(extents) - 1):
            if extents[pos][1] > extents[pos+1][0]:
                #TODO: Check the preceding line for an off-by-one error.
                raise ValueError("Overlapping extents")

        while extents:
            pos = infile.tell()
            if pos < extents[0][0]:
                # Copy non-audio data verbatim.
                temp = infile.read(min(CHUNK_SIZE, extents[0][0] - infile.tell()))
                outfile.write(temp)
            elif pos >= extents[0][0] and pos < extents[0][1]:
                # Byte-swap audio data while copying
                #TODO: Check for off-by-one error here.
                chunk = infile.read(min(CHUNK_SIZE, extents[0][1] - pos))
                assert len(chunk) % 2 == 0, "len(chunk) % 2 != 0:"

                outchunk = ''
                for x in range(0, len(chunk) / 2):
                    outchunk += chunk[(x*2) + 1] + chunk[x*2]
                outfile.write(outchunk)
            elif pos == extents[0][1]:
                # Reached the end of the current extent.
                extents.pop(0)
            else:
                # This should catch byte-swapping passing the end of the extent.
                raise Exception("Overshot our mark!")

            # Impose a post-condition on chunk copying and byte-swapping
            assert infile.tell() == outfile.tell(), "infile.tell() != outfile.tell()"

        while True:
            # Copy over any trailing data track without altering it.
            # (eg. a CD Extra data track)
            temp = infile.read(CHUNK_SIZE)
            if temp:
                outfile.write(temp)
            else:
                break

        # Close these before checking the post-condition
        infile.close()
        outfile.close()

        # Import a post-condition on the entire operation
        if os.stat(inpath).st_size != os.stat(outpath).st_size:
            # This indicates a bug in the program somewhere.
            raise IOError("Output filesize does not match input filesize!")
    except:
        if os.path.exists(outfile):
            # Make the conversion atomic
            os.remove(outfile)
        raise

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print "Usage: %s <cue file> <target bin file>" % sys.argv[0]
        print
        print "The source bin file will be automatically discovered by parsing the cue file."
        sys.exit(1)

    cuepath = os.path.abspath(sys.argv[1])

    # Get the AUDIO extents for the BIN file from the CUE file.
    binpath, track_extents = get_extents(cuepath)

    # Trying to have the same source and destination is too dangerous.
    if os.path.realpath(binpath) == os.path.realpath(sys.argv[2]):
        print ("Please specify a target filename different from the source"
               "filename contained in the CUE file.")
        sys.exit(2)

    swap_copy(binpath, sys.argv[2], track_extents)