#!/usr/bin/env python

version = '1.0'

import getopt, os, re, sys

def show_help():
    print 'Splits a dual layer DVD into two for burning to DVD+/-R(W)'
    print 'Usage: dvdsplit [OPTIONS]'
    print '  -1\t\t\tOnly create first disc'
    print '  -2\t\t\tOnly create second disc'
    print '  -a, --aid\t\tAudio track id (implies --ripper=mencoder)'
    print '  -c, --chapters\tJust show source chapter info'
    print '  -f, --format\t\tShow chapters in dvdauthor format'
    print '  -h, -?, --help\tShow this help message'
    print '  -p, --ripper\t\tRipping program:'
    print '\t\t\tstreamdvd: default'
    print '\t\t\tmencoder:  may discard extra angles and audio tracks'
    print '\t\t\tmplayer:   may crash'
    print '  -r, --range\t\tCreate single disc from range of chapters x-y'
    print '  -s, --source\t\tSource DVD device or file structure'
    print '  -t, --track\t\tSelect track (aka title)'
    print '  -v, --verbose\t\tShow extra information'
    print '  --title1\t\tTitle (name) for disc 1'
    print '  --title2\t\tTitle (name) for disc 2'
    print '  --titlebg1\t\tTitle background picture file for disc 1'
    print '  --titlebg2\t\tTitle background picture file for disc 2'
    print '  --menubg1\t\tMenu background picture file for disc 1'
    print '  --menubg2\t\tMenu background picture file for disc 2'

# Convert "HH:MM:SS" into seconds...
def hms_to_secs(hms):
    h, m, s = hms.split(':')
    return int(h) * 3600 + int(m) * 60 + int(s)

# ...and back
def secs_to_hms(s):
    h = s // 3600
    m = (s % 3600) // 60
    s %= 60
    return "%02d:%02d:%02d" % (h, m, s)

# Return a list of chapter lengths in form: "HH:MM:SS"
def get_chaplens(source, track):
    chaplens = []
    matchstr = \
            re.compile(r'\s*Chapter:\s+\d\d,\s+Length:\s+(\d\d:\d\d:\d\d)')
    lsdvd = os.popen('lsdvd -c -t %d %s' % (track, source), 'r')
    while True:
        line = lsdvd.readline()
        if not line: break
        match = matchstr.match(line)
        if match:
            chaplens.append(match.expand(r'\1'))
    lsdvd.close()
    return chaplens

# Returns comma-separated string
def get_langs(source, track):
    langs = None
    matchstr = re.compile(r'\s+Audio: \d+, Language: ([^\s]+) ')
    lsdvd = os.popen('lsdvd -a -t %d %s' % (track, source), 'r')
    while True:
        line = lsdvd.readline()
        if not line: break
        match = matchstr.match(line)
        if match:
            if langs:
                langs += ',' + match.expand(r'\1')
            else:
                langs = match.expand(r'\1')
    lsdvd.close()
    return langs

# Convert a list of chapter lengths to chapter start times
def chaplens_to_chaptimes(chaplens):
    chaptimes = []
    accsecs = 0
    for chap in chaplens:
        chaptimes.append(secs_to_hms(accsecs))
        accsecs += hms_to_secs(chap)
    return chaptimes

def print_chapinfo(chaptimes, chaplens):
    for n in range(len(chaptimes)):
        print "  Chapter %02d: %s + %s" % (n + 1, chaptimes[n], chaplens[n])

def run_streamdvd(source, track, first_chap, end_chap,
        verbose, filename):
    cmd = 'streamdvd -i %s -t %d -c %d-%d' % \
            (source, track, first_chap + 1, end_chap)
    if verbose:
        print 'Ripping to ' + filename + ' with ' + cmd
    outfile = open(filename, 'w')
    streamdvd = os.popen(cmd, 'r')
    while True:
        buf = streamdvd.read(8192)
        if not buf: break
        outfile.write(buf)

def run_mplayer(source, track, first_chap, end_chap,
        verbose, filename):
    cmd = ('mplayer', '-dumpstream', '-dumpfile', filename,
            '-dvd-device', source, '-chapter',
            '%d-%d' % (first_chap + 1, end_chap), 'dvd://%d' % track)
    if verbose:
        spc = ' '
        print 'Ripping with ' + spc.join(cmd)
    result = os.spawnvp(os.P_WAIT, cmd[0], cmd)
    if result:
        print >> sys.stderr, 'MPlayer rip failed'
        sys.exit(result)

def run_mencoder(source, track, first_chap, end_chap,
        verbose, filename, aid):
    cmd = ['mencoder', '-dvd-device', source, '-o', filename,
            '-of', 'mpeg', '-oac', 'copy', '-ovc', 'copy',
            '-mpegopts', 'format=dvd',
            '-chapter', '%d-%d' % (first_chap + 1, end_chap)]
    if not aid:
        cmd.append('-aid')
        cmd.append(aid)
    cmd.append('dvd://%d' % track)
    if verbose:
        spc = ' '
        print 'Ripping with ' + spc.join(cmd)
    result = os.spawnvp(os.P_WAIT, cmd[0], cmd)
    if result:
        print >> sys.stderr, 'MEncoder rip failed'
        sys.exit(result)

def run_wizard(infile, disc, title, langs, chaptimes, verbose,
        titlebg, menubg):
    outimg = 'disc%d' % disc
    cmd = ['dvdwizard', '-T', title, '-o', outimg, '-A', langs]
    if len(chaptimes) > 1:
        chapstring = chaptimes[0]
        for chap in chaptimes[1:]:
            chapstring += ',' + chap
        cmd.append('-c')
        cmd.append(chapstring)
    if titlebg:
        cmd.append('-B')
        cmd.append(titlebg)
    if menubg:
        cmd.append('-b')
        cmd.append(menubg)
    cmd.append(infile)
    if verbose:
        spc = ' '
        print 'Generating ' + outimg + ' with ' + spc.join(cmd)
    result = os.spawnvp(os.P_WAIT, cmd[0], cmd)
    if result:
        print >> sys.stderr, 'dvdwizard failed'
        sys.exit(result)

# Note chapter numbering starts from 0 and end_chap is exclusive
def dvd_process(source, track, chaplens, title, titlebg, menubg,
        first_chap, end_chap, verbose, disc, ripper, aid, langs):
    filename = 'ripped%d.mpg' % disc
    replex = False
    if ripper == 'streamdvd':
        run_streamdvd(source, track, first_chap, end_chap,
                verbose, filename)
    elif ripper == 'mplayer':
        run_mplayer(source, track, first_chap, end_chap,
                verbose, filename)
    else:
        run_mencoder(source, track, first_chap, end_chap,
                verbose, filename, aid)
        replex = True
    
    if replex:
        if not aid:
            aid = '0x80'
        cmd = ('replex', '-k', '-i', 'PS', '-c', aid, '-t', 'DVD',
                '-o', 'replexed.mpg', filename)
        if verbose:
            spc = ' '
            print 'Remultiplexing with ' + spc.join(cmd)
        result = os.spawnvp(os.P_WAIT, cmd[0], cmd)
        if result:
            print >> sys.stderr, 'replex failed'
            sys.exit(result)
        os.remove(filename)
        filename = 'replexed.mpg'

    run_wizard(filename, disc, title, langs,
            chaplens_to_chaptimes(chaplens[first_chap:end_chap]),
            verbose, titlebg, menubg)

    os.remove(filename)

def main():
    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], '12a:c?fh:p:qr:s:t:',
                ('aid=', 'chapters', 'format', 'help', 'quiet', 'source=',
                'ripper=', 'range=', 'track=', 'title1=', 'title2=',
                'titlebg1=', 'titlebg2=', 'menubg1=', 'menubg2='))
    except getopt.GetoptError:
        show_help()
        sys.exit(1)

    source = '/dev/dvd'
    chaps_only = False
    verbose = True
    first_chap = 0
    last_chap = 0
    track = 1
    only_disc1 = False
    only_disc2 = False
    ripper = 'streamdvd'
    aid = None
    title1 = None
    title2 = None
    titlebg1 = None
    titlebg2 = None
    menubg1 = None
    menubg2 = None
    format_chaps = False
    
    for o, a in opts:
        if o == '-1':
            only_disc1 = True
        elif o == '-2':
            only_disc2 = True
        elif o in ('-a', '--aid'):
            aid = a
        elif o in ('-c', '--chapters'):
            chaps_only = True
        elif o in ('-f', '--format'):
            format_chaps = True
        elif o in ('-h', '-?', '--help'):
            show_help()
            sys.exit(0)
        elif o in ('-q', '--quiet'):
            verbose = False
        elif o in ('-p', '--ripper'):
            if a != 'streamdvd' and a != 'mencoder' and a != 'mplayer':
                print >> sys.stderr, "Bad ripper option"
                sys.exit(1)
            ripper = a
        elif o in ('-r', '--range'):
            try:
                crange = a.split('-')
                assert len(crange) == 2
                assert(crange[0] > 0 and crange[1] >= crange[0])
                first_chap = int(crange[0])
                last_chap = int(crange[1])
                only_disc1 = True
            except:
                print >> sys.stderr, "Bad range option"
                sys.exit(1)
        elif o in ('-s', '--source'):
            source = a
        elif o in ('-t', '--track'):
            try:
                track = int(a)
            except:
                print >> sys.stderr, "Bad track option"
                sys.exit(1)
        elif o == '--title1':
            title1 = a
        elif o == '--title2':
            title2 = a
        elif o == '--titlebg1':
            titlebg1 = a
        elif o == '--titlebg2':
            titlebg2 = a
        elif o == '--menubg1':
            menubg1 = a
        elif o == '--menubg2':
            menubg2 = a
    
    if (only_disc2 and first_chap):
        print >> sys.stderr, "Overriding -2 option because chapter range given"
        only_disc2 = False
    
    chaplens = get_chaplens(source, track)
    chaptimes = chaplens_to_chaptimes(chaplens)

    if verbose or chaps_only or format_chaps:
        if format_chaps:
            chapstr = ''
            for n in range(len(chaptimes)):
                if chapstr:
                    chapstr += ','
                chapstr += chaptimes[n]
            print chapstr
        else:
            print "Source chapters"
            print_chapinfo(chaptimes, chaplens)
        if chaps_only:
            return 0
    
    langs = get_langs(source, track)
    if not langs:
        print >> sys.stderr, "Can't find audio track languages"
        sys.exit(1)
    if verbose:
        print 'Audio track languages: ' + langs
    
    if (not title1 and not only_disc2) or (not title2 and not only_disc1):
        print >> sys.stderr, "Missing title"
        sys.exit(1)

    if titlebg1 and not menubg1:
        menubg1 = titlebg1
    elif menubg1 and not titlebg1:
        titlebg1 = menubg1
    if titlebg2 and not menubg2:
        menubg2 = titlebg2
    elif menubg2 and not titlebg2:
        titlebg2 = menubg2
    if titlebg1 and not titlebg2:
        titlebg2 = titlebg1
    if titlebg2 and not titlebg1:
        titlebg1 = titlebg2
    if menubg1 and not menubg2:
        menubg2 = menubg1
    if menubg2 and not menubg1:
        menubg1 = menubg2

    if first_chap:
        chaplens = chaplens[first_chap - 1:last_chap - 1]
        chaptimes = chaplens_to_chaptimes(chaplens)
        split_chap = last_chap - 1
        first_chap -= 1
    else:
        # Find chapter nearest halfway point
        total_time = hms_to_secs(chaptimes[-1]) + hms_to_secs(chaplens[-1])
        if verbose:
            print "Total time: " + secs_to_hms(total_time)
        half_time = total_time / 2
        nearest_diff = 1000000
        split_chap = 0
        for n in range(len(chaptimes)):
            chap_secs = hms_to_secs(chaptimes[n])
            chap_diff = abs(chap_secs - half_time)
            if chap_diff < nearest_diff:
                nearest_diff = chap_diff
                split_chap = n
        if split_chap <= 1:
            print >> sys.stderr, "Unable to split this DVD (no chapters?)"
            sys.exit(1)
        if verbose:
            print "Splitting at chapter %02d (%s)" % \
                    (split_chap + 1, chaptimes[split_chap])

    if not only_disc2:
        dvd_process(source, track, chaplens, title1, titlebg1, menubg1,
                first_chap, split_chap, verbose, 1, ripper, aid, langs)
    if not only_disc1:
        dvd_process(source, track, chaplens, title2, titlebg2, menubg2,
                split_chap, len(chaplens), verbose, 2, ripper,
                        aid, langs)
    
    return 0

if __name__ == '__main__':
    main()