#!/usr/bin/env python

from lxml import etree
import shutil
import signal
import subprocess
import urllib


BASE_URL = "http://player.sbs.com.au"

def grab_xml(path):
	f = urllib.urlopen(BASE_URL + path)
	doc = etree.parse(f)
	f.close()
	return doc

def choose(options, allow_multi):
	skeys = sorted(options.keys())
	for i, key in enumerate(skeys):
		print " %d) %s" % (i+1, key)
	print " 0) Back"
	while True:
		try:
			values = map(int, raw_input("Choose> ").split())
			if len(values) == 0:
				continue
			if 0 in values:
				return
			values = [options[skeys[value-1]] for value in values]
			if allow_multi:
				return values
			else:
				if len(values) == 1:
					return values[0]
		except ValueError:
			pass

def get_menu_list():
	settings = grab_xml("/playerassets/programs/config/standalone_settings.xml")
	menu_url = settings.xpath("/settings/setting[@name='menuURL']/@value")[0]

	root_menu = grab_xml(menu_url)
	menu_list = {}
	for menu in root_menu.xpath("//menu"):
		try:
			title = menu.xpath("title/text()")[0]
			playlist_url = menu.xpath("playlist/@xmlSrc")[0]
			if title in menu_list:
				i = 2
				while True:
					if (title+str(i)) not in menu_list:
						title += str(i)
						break
					i += 1
			menu_list[title] = playlist_url
		except IndexError:
			continue
	
	return menu_list

def get_video_list(playlist_url):
	video_list = {}
	playlist = grab_xml(playlist_url)
	for video_desc in playlist.xpath("//video"):
		video_desc_url = video_desc.xpath("@src")[0]
		title = video_desc.xpath("title/text()")[0].strip()
		video_list[title] = (title, video_desc_url)
	return video_list

def download_video(title, video_desc_url):
	video = grab_xml(video_desc_url)
	vbase = video.xpath("//meta/@base")[0]
	bestrate = 0
	bestvpath = None
	for vpath in video.xpath("//switch/video"):
		rate = float(vpath.xpath("@system-bitrate")[0])
		if rate > bestrate:
			bestrate = rate
			bestvpath = vpath.xpath("@src")[0]
	filename = title + "." + bestvpath.rsplit(".", 1)[1]
	if vbase.startswith("rtmp://"):
		download_rtmp(filename, vbase, bestvpath)
	else:
		download_urllib(filename, vbase + bestvpath)

def download_rtmp(filename, vbase, vpath):
	if vpath.endswith(".flv"):
		vpath = vpath[:-4]
	cmd = [
		"rtmpdump",
		"-o", filename,
		"-r", vbase,
		"-y", vpath,
	]
	try:
		p = subprocess.Popen(cmd)
		p.wait()
	except KeyboardInterrupt:
		print "Cancelled", cmd
		try:
			p.terminate()
			p.wait()
		except KeyboardInterrupt:
			p.send_signal(signal.SIGKILL)
			p.wait()

def download_urllib(filename, url):
	print "Downloading: %s -> %s" % (url, filename)
	try:
		src = urllib.urlopen(url)
		dst = open(filename, "w")
		shutil.copyfileobj(src, dst)
	except KeyboardInterrupt:
		print "\nCancelled", url
	finally:
		src.close()
		dst.close()

def main():
	while True:
		menu_list = get_menu_list()
		playlist_url = choose(menu_list, allow_multi=False)
		if playlist_url is None:
			continue
		video_list = get_video_list(playlist_url)
		while True:
			generator = choose(video_list, allow_multi=True)
			if generator is None:
				break
			for title, video_desc_url in generator:
				download_video(title, video_desc_url)

if __name__ == "__main__":
	try:
		main()
	except (KeyboardInterrupt, EOFError):
		print "\nExiting..."