#!/usr/bin/env python

"""
dvdmenu
Creates a menu file for dvdauthor
(c) 2007 Tony Houghton
Home page: <http://www.realh.co.uk>
Distributed under the GPL: <http://www.gnu.org/licenses/gpl.html>
"""

version = (0, 0, 1)

import math
import optparse
import os
import select
import subprocess
import sys

SPU_PNG_NAMES = ('spu_normal.png', 'spu_select.png', 'spu_highlight.png')
MPEG_FORMAT = '8'

options = None
fps = 25
height = 576

def _(s): return s

def parse_opts():
	parser = optparse.OptionParser(version = "%d.%d.%d" % version,
			description = _('dvdmenu creates a menu file for dvdauthor. ' + \
					'The result is a file called "menu.mpg" in the ' + \
					'current directory. A number of temporary files are ' + \
					'also created there.'),
			prog = "dvdmenu")
	parser.add_option('-b', '--background', default = "",
			help = _("Image file to use as the background. May be " + \
			"ommitted to use plain black. Temporary files called " + \
			"background.png and composite.png are created in either case"))
	parser.add_option('-f', '--format', default = "button%02d.png",
			help = _('printf format string to use for buttons eg ' + \
			'"button%02d.png" which is the default. dvdmenu starts ' + \
			'counting from 0 and continues until there are no more ' + \
			'matching files. This determines how many buttons to generate'))
	parser.add_option('-r', '--rows', type = "int", default = 3,
			help = _("How many rows of buttons to use. If there are fewer " + \
			"buttons than rows, the latter is rounded down. This also " + \
			"determines the final size of each button. Default value 3"))
	parser.add_option('-m', '--margin', default = "50",
			help = _("A single value, or comma-separated values. " + \
			"The width in pixels of the margin at each edge. 4 values " + \
			"are top,right,bottom,left, 2 values are vert,horiz. " + \
			"The default is 50. Use even numbers."))
	parser.add_option('-p', '--spacing', type = "int", default = 20,
			help = _("The number of pixels between buttons. The default is " + \
			"20. Use even numbers."))
	parser.add_option('-w', '--widescreen', action = "store_true",
			help = _("Make a 16:9 anamorphic image instead of 4:3. " + \
			"NB dvdmenu assumes all the input image files have the " + \
			"appropriate aspect ratio regardless of their dimensions"))
	parser.add_option('-n', '--ntsc', action = "store_true",
			help = _("Use NTSC format. The default is PAL."))
	parser.add_option('-s', '--sound', default = "",
			help = _("Specify a sound file " + \
			"(MPEG-1 layer 2 audio) to use instead of silence. " + \
			"Otherwise temporary files silence.wav and silence.mpa " + \
			"will be created. " + \
			"Consider using in conjunction with --frames to make the " + \
			"picture last the same time as the sound"))
	parser.add_option('-a', '--frames', type = "int", default = 1,
			help = _("How many frames the picture should last"))
	parser.add_option('-t', '--thickness', type = "int", default = 2,
			help = _("Thickness of button borders in pixels (default 2). " + \
			"Even numbers only."))
	parser.add_option('-c', '--colours', default = "#000000:#40c0ff:#c0f0ff",
			help = _('Colon-separated list of 3 colours to use for button ' + \
			'borders for inactive, selected and highlighted buttons ' + \
			'respectively. Each colour is a string recognised by ' + \
			'ImageMagick. The default is "#0000c0:#40c0ff:#c0f0ff"'))
	parser.add_option('-e', '--transparent', default = '#000000',
			help = _('A colour instead of black ("#000000") to convert ' + \
			'to transparency in case you want to use black borders with ' + \
			'the --colours option.'))
	global options
	global fps, height
	options = parser.parse_args(sys.argv)[0]
	if options.ntsc:
		fps = 30000/1001
		height = 480
	margin = options.margin.split(',')
	for i in range(len(margin)):
		margin[i] = int(margin[i]) & ~1
	if len(margin) == 2:
		margin = [margin[0], margin[1], margin[0], margin[1]]
	elif len(margin) == 1:
		margin = [margin[0], margin[0], margin[0], margin[0]]
	elif len(margin) != 4:
		print >>sys.stderr, _("Invalid margin parameter")
		sys.exit(1)
	options.margin = margin
	if options.spacing & 1:
		options.spacing += 1
	if options.spacing < 2:
		options.spacing = 2
	if options.thickness & 1:
		options.thickness += 1
	options.colours = options.colours.split(':')


def find_in_path(leafnames, doesnt_matter = False):
	if isinstance(leafnames, str):
		leafnames = [leafnames]
	for leafname in leafnames:
		for d in os.environ['PATH'].split(':'):
			f = os.path.join(d, leafname)
			if os.path.exists(f) and os.access(f, os.X_OK):
				break
		else:
			if not doesnt_matter:
				print >>sys.stderr, _("Can't find %s in PATH") % leafname
				sys.exit(1)
			return False
	return True


def generate_silence(length):
	BS = 1024
	aframes = int(length * 48000)
	soxcmd = 'sox -r 48000 -c 1 -t .sb - -c 1 silence.wav'
	sox = subprocess.Popen(soxcmd.split(), stdin = subprocess.PIPE)
	while aframes:
		dsize = min(BS, aframes)
		data = "\x00" * dsize
		aframes -= dsize
		sox.stdin.write(data)
	sox.stdin.close()
	sox.wait()
	subprocess.call('mp2enc -b 128 -r 48000 -s -o silence.mpa < silence.wav',
			shell = True)
	os.unlink('silence.wav')


def generate_black_background():
	subprocess.call(['convert', '-size', '720x%d' % height, '-depth', '16',
		'xc:black', 'background.png'])


def scale_background(filename):
	subprocess.call(['convert', '-resize', '720x%d!' % height, '-depth', '16',
		filename, '-background', 'black', 'background.png'])


def scan_for_buttons(format):
	buttons = []
	n = 0
	while os.path.isfile(format % n):
		buttons.append(format % n)
		n += 1
	if not buttons:
		print >>sys.stderr, _("No button files. Check your --format option.")
		sys.exit(1)
	return buttons


def calculate_button_size(columns, rows):
	# '& ~1' rounds down to nearest even number
	aspect = 720.0 / float(height)
	max_w = int((720 - options.margin[1] - options.margin[3] - \
			(columns - 1) * options.spacing) / columns) & ~1
	max_h = int((height - options.margin[0] - options.margin[2] - \
			(rows - 1) * options.spacing) / rows) & ~1
	if float(max_w) / float(max_h) > aspect:
		max_w = int(max_h * aspect) & ~1
	else:
		max_h = int(max_w / aspect) & ~1
	return [max_w, max_h]


def layout_buttons(columns, rows, n_buttons, size, margins, spacing):
	positions = []
	n = 0
	extra_h_margin = int((720 - margins[1] - margins[3] - columns * size[0] - \
			(columns - 1) * spacing) / 2) & ~1
	extra_v_margin = int((height - margins[0] - margins[2] - rows * size[1] - \
			(rows - 1) * spacing) / 2) & ~1
	for y in range(rows):
		ypos = margins[0] + extra_v_margin + y * (size[1] + spacing)
		for x in range(columns):
			xpos = margins[3] + extra_h_margin + x * (size[0] + spacing)
			positions.append([xpos, ypos])
			n += 1
			if n >= n_buttons:
				return positions
	print >>sys.stderr, _("BUG: More buttons than rows * columns")
	sys.exit(1)


def generate_composite(button_files, button_positions, button_size):
	cmd = ['convert', 'background.png']
	for n in range(len(button_files)):
		cmd += ['-draw', 'image src-over %d,%d %d,%d %s' % \
				(button_positions[n][0], button_positions[n][1],
				button_size[0], button_size[1], button_files[n])]
	subprocess.call(cmd + ['composite.png'])


def generate_png(filename, button_positions, button_size,
		colour, transparent, width):
	# colour is string in any format recognised by Image Magick
	cmd = ['convert', '-size', '720x%d' % height,
			'-fill', transparent, '-stroke', colour, '-strokewidth', str(width)]
	for b in button_positions:
		cmd += ['-draw',
				'rectangle %d,%d %d,%d' % (b[0], b[1],
						b[0] + button_size[0],
						b[1] + button_size[1])]
	# For some reason background image has to come after draw primitives
	# otherwise it obliterates them
	subprocess.call(cmd + ['xc:' + transparent, filename])
	subprocess.call(['mogrify', '-colors', '4', '-matte',
			'-type', 'PaletteMatte', '-transparent', transparent, filename])


def generate_spu_pngs(button_positions, button_size):
	for n in range(3):
		generate_png(SPU_PNG_NAMES[n], button_positions, button_size,
				options.colours[n], options.transparent, options.thickness)


#def generate_spu_xml(button_positions, button_size):
#	fp = open("spumux.xml", 'w')
#	fp.write("""<subpictures>
#	<stream>
#		<spu
#			force="yes"
#			start="00:00:00.00"
#			image=""" + '"' + os.path.abspath(SPU_PNG_NAMES[0]) + '"' + """
#			select=""" + '"' + os.path.abspath(SPU_PNG_NAMES[1]) + '"' + """
#			highlight=""" + '"' + os.path.abspath(SPU_PNG_NAMES[2]) + '"' + """
#			autooutline="infer"
#			outlinewidth=""" + '"' + str(options.spacing - 1) + '"' + """
#			autoorder="rows"
#		/>
#	</stream>
#</subpictures>""")
#	fp.close()


def generate_spu_xml(button_positions, button_size):
	# Autodetection of buttons isn't working properly, perhaps because of
	# convert stupidly adding extra colours
	fp = open("spumux.xml", 'w')
	fp.write("""<subpictures>
	<stream>
		<spu
			force="yes"
			start="00:00:00.00"
			image=""" + '"' + os.path.abspath(SPU_PNG_NAMES[0]) + '"' + """
			select=""" + '"' + os.path.abspath(SPU_PNG_NAMES[1]) + '"' + """
			highlight=""" + '"' + os.path.abspath(SPU_PNG_NAMES[2]) + '" >')
	for b in button_positions:
		fp.write("""
			<button
				x0=""" + '"' + str(b[0]) + '"' + """
				y0=""" + '"' + str(b[1]) + '"' + """
				x1=""" + '"' + str(b[0] + button_size[0]) + '"' + """
				y1=""" + '"' + str(b[1] + button_size[1]) + '" />')
	fp.write("""
		</spu>
	</stream>
</subpictures>""")
	fp.close()


def convert_background_to_mpeg():
	png2yuv = subprocess.Popen(['png2yuv', '-j', 'composite.png',
			'-f', str(fps), '-I', 'p', '-n', str(options.frames)],
			stdout = subprocess.PIPE)
	if options.ntsc:
		fr_code = '4'
		tvsys = 'p'
	else:
		fr_code = '3'
		tvsys = 'n'
	if options.widescreen:
		aspect_code = '3'
	else:
		aspect_code = '2'
	mpeg2enc = subprocess.Popen(['mpeg2enc', '-F', fr_code, '-f', MPEG_FORMAT,
		'-b', '7400', '--cbr', '-n', tvsys, '-a', aspect_code,
		'-o', 'background.m2v'], stdin = png2yuv.stdout)
	png2yuv.wait()
	mpeg2enc.wait()


def multiplex_background(sound_file):
	fp = open('menu.mpg', 'w')
	mplex = subprocess.Popen(['mplex', '-f', MPEG_FORMAT,
			'-o', '/dev/stdout', 'background.m2v', sound_file],
			stdout = subprocess.PIPE)
	spumux = subprocess.Popen(['spumux', 'spumux.xml'], stdin = mplex.stdout,
			stdout = fp)
	mplex.wait()
	spumux.wait()
	fp.close()


def generate_menu():
	global options

	find_in_path(['png2yuv', 'mpeg2enc'])

	if not options.sound:
		find_in_path(['mp2enc', 'sox'])

	if not options.background:
		find_in_path('convert')
		generate_black_background()
	else:
		scale_background(options.background)
	
	button_files = scan_for_buttons(options.format)
	n_buttons = len(button_files)
	if n_buttons < options.rows:
		options.rows = n_buttons
	n_columns = int(math.ceil(float(n_buttons) / float(options.rows)))
	button_size = calculate_button_size(n_columns, options.rows)
	button_positions = layout_buttons(n_columns, options.rows, n_buttons,
			button_size, options.margin, options.spacing)
	generate_composite(button_files, button_positions, button_size)
	os.unlink('background.png')

	generate_spu_pngs(button_positions, button_size)
	generate_spu_xml(button_positions, button_size)

	if options.sound:
		sound_file = options.sound
	else:
		generate_silence(float(options.frames) / fps)
		sound_file = 'silence.mpa'

	convert_background_to_mpeg()
	os.unlink('composite.png')
	multiplex_background(sound_file)
	if not options.sound:
		os.unlink('silence.mpa')
	os.unlink('background.m2v')
	os.unlink('spumux.xml')
	for n in range(3):
		os.unlink(SPU_PNG_NAMES[n])


if __name__ == '__main__':
	parse_opts()
	generate_menu()