#!/usr/bin/env python

import os, sys, math, glob, re

vobname = sys.argv[1]

# use tcprobe to see the list of audio channels:
#os.system("tcprobe -i '%s'" % vobname)

# list audio channels to be transcoded (ATM, don't know how to get
# ac3/mp3 information):
#audioList = [("mp2", 64, "de")]
audioList = [("ac3", 0, "de")]
#audioList = [("ac3", 0, "de"), ("ac3", 1, "en")]
videoPTSOffset = ""
#videoPTSOffset = "-O -336ms"
# 0x1
#  PCM uncompressed audio
# 0x50
#  MPEG layer-2 aka MP2
# 0x55
#  MPEG layer-3 aka MP3. Also have a look at --lame_preset if you intend to do VBR audio.
# 0x2000
#  AC3 audio
# 0xfffe
#  OGG/Vorbis audio

# --------------------------------------------------------------------

basename, ext = os.path.splitext(vobname)
_, basename = os.path.split(basename)

m2vname = basename + ".m2v"

def verboseUnlink(fn):
	print "rm '%s'" % fn
	if os.unlink(fn):
		print "# unlink failed!"

def command(cmd, output = None):
	print cmd
	if output and len(glob.glob(output)):
		print "# skipped, because '%s' already exists.." % (output, )
		return
	rc = os.system(cmd)
	ec = 0
	if rc >= 256:
		print "# COMMAND TERMINATED WITH EXIT CODE %d!" % (rc / 256)
		ec = 1
	elif rc > 0:
		print "# COMMAND TERMINATED BY SIGNAL %d!" % rc
		ec = 2
	if ec:
		if output and os.path.exists(output):
			verboseUnlink(output)
		sys.exit(ec)

audioFiles = []

if False: # use_transcode
	# input range for testing (video frame/audio byte - not working
	# like this :-( )
	
	# extract video stream with Transcode
	command("tcextract -i '%s' -t vob -x mpeg2 > '%s'" %
			(vobname, m2vname), m2vname)
	
	# dump selected audio streams with Transcode
	for audioType, audioID, audioLang in audioList:
		fn = "%s_%d.%s" % (basename, audioID, audioType)
		command("tcextract -i '%s' -t vob -x %s -a %d > '%s'" %
				(vobname, audioType, audioID, fn), fn)
		audioFiles.append(fn)
elif False:
	# extract video stream with MPlayer
	command("mplayer -dumpvideo -dumpfile '%s' '%s'" %
			(m2vname, vobname), m2vname)

	# dump selected audio streams with MPlayer
	for audioType, audioID, audioLang in audioList:
		fn = "%s_%d.%s" % (basename, audioID, audioType)
		command("mplayer -dumpaudio -dumpfile '%s' -aid %d '%s'" %
				(fn, 128+audioID, vobname), fn)
		audioFiles.append(fn)
else:
	projectXLog = basename + "_log.txt"
	command("projectx -demux '%s'" % (vobname, ), projectXLog)

	# FIXME: state-based recognition of:
	# --> AC-3/DTS Audio (SubID 0x82)
	# ---> new File: '/media/Hitch/stream[2].ac3'
	# for correct audioIDs!

# 	re_projectXAudio = re.compile(
# 		r"^Audio ([0-9]+) \(([a-z0-9]+)\):	.*'([^']+)'$", re.MULTILINE)

# 	pxAudioFiles = {}
# 	for ma in re_projectXAudio.finditer(file(projectXLog).read()):
# 		pxAudioFiles[int(ma.group(1))] = (ma.group(2), ma.group(3))

	re_projectXAudio = re.compile(
		r"^--> .* Audio \(?(?:SubID |on PID )?(0x[0-9A-F]+)\)?|^---> new File: '(.*\.)(.*?)'$",
		re.MULTILINE)

	pxAudioFiles = {}
	aid = None
	for ma in re_projectXAudio.finditer(file(projectXLog).read()):
		if ma.group(1):
			aid = eval(ma.group(1)) - 128
		else:
			if aid != None:
				pxAudioFiles[aid] = (ma.group(3), ma.group(2)+ma.group(3))
				aid = None

	print "found AIDs:", pxAudioFiles.keys()

	good = True

	if not os.path.exists(m2vname):
		good = False
		sys.stderr.write("Make sure that '%s' exists!\n" % m2vname)

	for audioType, audioID, audioLang in audioList:
		fn = pxAudioFiles[audioID][1]
		realType = pxAudioFiles[audioID][0]
		if audioType != realType:
			sys.stderr.write("WARNING: %s has type %s, not %s..\n" % (
				fn, realType, audioType))

		#fn = "%s_%d.%s" % (basename, audioID, audioType)
		if not os.path.exists(fn):
			good = False
			sys.stderr.write("Make sure that '%s' exists!\n" % fn)
		else:
			print "'%s' found. :-)" % fn
		audioFiles.append(fn)

	assert good

# convert .ac3 -> .mp2
if False:
  for i, fn in enumerate(audioFiles):
	audiobase, ext = os.path.splitext(fn)
	if ext == ".ac3" and not os.path.exists(audiobase + ".mp2"):
		# HACK: increase volume by 6dB (usually too low IMHO)
		audioFilter = " -af volume=6"
		command("mplayer%s -ao pcm:waveheader:file='%s' -vc null -vo null '%s'"
				% (audioFilter, audiobase + ".wav", fn), audiobase + ".wav")
		command("toolame -b 224 '%s' '%s'"
				% (audiobase + ".wav", audiobase + ".mp2"), audiobase + ".mp2")
		verboseUnlink(audiobase + ".wav")
		audioFiles[i] = audiobase + ".mp2"

Mb = 1024*1024
Gb = 1024*Mb

# calculate shrink factor
available  = 4700000000 # 4,7"GB", actually around 4.37 Gb
available -= 112*Mb # reserved space (add more for subtitles etc.)

for fn in audioFiles:
	print "# audio has %d Mb: %s" % ((os.path.getsize(fn)+Mb-1)/Mb, fn)
	available -= os.path.getsize(fn)

factor = math.ceil(os.path.getsize(m2vname) * 1000.0 / available) / 1000

print "# shrink factor: %s (%s -> %s Mb)" % (
	factor, os.path.getsize(m2vname)/Mb, int(available/Mb))

# shrink if necessary
if factor > 1.0:
	command("tcrequant -i '%s' -o shrinked.m2v -f %s" %
			(m2vname, factor))
	os.rename("shrinked.m2v", m2vname)
else:
	print "# (skipped tcrequant)"

mplexed = basename + "_%d_dvd47.mpeg"
mplexedGlob = mplexed.replace("%d", "[0-9]")

# re-multiplexing:
command("mplex -f 8 %s -o '%s' '%s' " % (videoPTSOffset, mplexed, m2vname) +
		" ".join(["'%s'" % fn for fn in audioFiles]), mplexedGlob)

for fn in [m2vname] + audioFiles:
	os.unlink(fn)

daXML = """<dvdauthor>
  <vmgm />
  <titleset>
    <titles>
"""
for _, _, lang in audioList:
	daXML += "      <audio lang=\"%s\" />\n" % lang
daXML += "      <pgc>\n"
for fn in glob.glob(mplexedGlob):
	daXML += "        <vob file=\"%s\"/>\n" % fn
daXML += """      </pgc>
    </titles>
  </titleset>
</dvdauthor>
"""

xmlFilename = basename + ".xml"
print "# writing XML file '%s'..." % (xmlFilename, ),
if os.path.exists(xmlFilename):
	print "skipped (existing)."
else:
	file(xmlFilename, "w").write(daXML)
	print "done."

burndir = basename + ".burndir"
command("dvdauthor -o '%s' -x '%s'" % (burndir, xmlFilename), burndir)

print "Done. Use e.g. one of the following for testing:"
print "mplayer -dvd-device '%s' dvd://" % (burndir, )
print "xine '%s'" % (burndir, )