#!/usr/bin/env python
'''
Usage: splitpea MODE [OPTIONS] FILE

Modes:
  -s, --split            Splits the file specifed by FILE, naming each
                         piece FILE.001, FILE.002, ...
  -j, --join             Joins pieces named FILE.001, FILE.002, ... into
                         a single file named FILE.

Split Mode Options:
  -b SIZE, --bytes=SIZE  Outputs pieces of SIZE bytes, default is 1m.

SIZE may have a multiplier suffix: b for 512 bytes, k for kilobytes, m for
megabytes.

General Options:
  -?, -h, --help         Show this help.
  --version              Show the program version.
  -v, --verbose          Displays more information while running
'''

__title__ = 'splitpea'
__author__ = 'Gordon Tyler'
__version__ = '1.0'

import sys
import os.path
import getopt
import string

IO_CHUNK_SIZE = 1048576

def printVersion():
    print '%s version %s' % (__title__, __version__)

def printHeader():
    printVersion()
    print 'Copyright (C) 2003, %s' % __author__
    print 'This may be free distributed under the terms of the GNU GPL.'
    print

def printUsage(errorMessage = ''):
    printHeader()

    if errorMessage != '':
        print 'Error: %s' % errorMessage

    print __doc__
    return

def printError(errorMessage):
    print '%s: %s' % (__title__, message)

def main(args):
    try:
        options, args = getopt.getopt(args, '?hsjb:v', ['help', 'split', 'join', 'bytes=', 'version', 'verbose'])
    except getopt.error, detail:
        printError(detail)
        return

    splitSize = 1048576
    mode = ''
    verbose = 0

    for option in options:
        if option[0] == '-?' or option[0] == '-h' or option[0] == '--help':
            printUsage()
            return
        elif option[0] == '--version':
            printVersion()
        elif option[0] == '-s' or option[0] == '--split':
            mode = 'split'
        elif option[0] == '-j' or option[0] == '--join':
            mode = 'join'
        elif option[0] == '-b' or option[0] == '--bytes':
            suffix = option[1][-1:]
            if string.count('bkm', suffix) > 0:
                splitSize = int(option[1][:-1])
                if suffix == 'b':
                    splitSize = splitSize * 512
                elif suffix == 'k':
                    splitSize = splitSize * 1024
                elif suffix == 'm':
                    splitSize = splitSize * 1048576
                else:
                    printUsage('Unrecognised size suffix.')
                    return
            else:
                splitSize = int(option[1])
        elif option[0] == '-v' or option[0] == '--verbose':
            verbose = 1
            
    if mode == '':
        printUsage('Must specify a mode')
        return
    
    if mode == 'split':
        if len(args) == 0:
            printUsage('Must specify a file to split')
            return

        if not os.path.exists(args[0]):
            printError('%s: No such file' % args[0])
            return
        
        inputFilename = args[0]
        
        try:
            input = open(inputFilename, 'rb')
        except IOError, detail:
            printError(detail)
            return

        outputCount = 1
        outputRemaining = 0
        output = None

        try:
            data = input.read(IO_CHUNK_SIZE)
            while len(data) > 0:
                if outputRemaining == 0:
                    if output:
                        output.close()
                        if verbose > 0:
                            print 'done.'
                        outputCount = outputCount + 1
                        
                    outputFilename = inputFilename + '.%03d' % outputCount
    
                    if verbose > 0:
                        print 'Writing piece: %s... ' % outputFilename,
    
                    output = open(outputFilename, 'wb')
                    outputRemaining = splitSize
                
                if outputRemaining >= len(data):
                    output.write(data)
                    outputRemaining -= len(data)
                    data = input.read(IO_CHUNK_SIZE)
                else:
                    output.write(data[:outputRemaining])
                    data = data[outputRemaining:]
                    outputRemaining = 0

            if output:
                output.close()
                if verbose > 0:
                    print 'done.'
        except IOError, detail:
            printError(detail)

        input.close()
    elif mode == 'join':
        if len(args) == 0:
            printUsage('Must specify the original split filename')
            return

        prefix = args[0]

        try:
            output = open(prefix, 'wb')
        except IOError, detail:
            printError(detail)
            return

        done = 0
        inputCount = 1

        while (not done):
            inputFilename = prefix + '.%03d' % inputCount
            if os.path.exists(inputFilename):
                try:
                    if verbose > 0:
                        print 'Reading piece: %s... ' % inputFilename,
                        
                    input = open(inputFilename, 'rb')
                    data = input.read(IO_CHUNK_SIZE)
                    while len(data) > 0:
                        output.write(data)
                        data = input.read(IO_CHUNK_SIZE)
                    input.close()

                    if verbose > 0:
                        print 'done.'
                except IOError, detail:
                    printError(detail)
                    return

                inputCount = inputCount + 1
            else:
                done = 1

        output.close()
        
    return

if __name__ == '__main__':
    main(sys.argv[1:])
