#!/usr/bin/python

# (C) by Robby Stephenson, http://periapsis.org, 2006
# GPL v2
# Script for searching and importing data from amarok into tellico
# http://periapsis.org/tellico/
# http://amarok.kde.org

import os
import sys
import string
import locale
import md5
import base64
import getopt

NAME = os.path.split(sys.argv[0])[1]
VERSION = "0.4"

options = {}

albumnames = {}
artistnames = {}
genrenames = {}
yearnames = {}
coverdir = ""
images = set()

def usage():
    sys.stderr.write(
"""This is %s version %s,
a tool for importing info from amaroK into Tellico.
Written by Robby Stephenson. Copyright (c) 2006.
http://periapsis.org

Usage: %s [-h] [-v] [-q] [-x] [-t title] [-a artist] [-k genre]

-t: Search for album or songtitle
-a: Search for artist
-k: Search for genre
-x: Do not include images
-q: Suppress output, useful for debugging
-h: Display this text
-v: Be verbose in reporting what is done

Multiple search terms are combined and assumed
to imply an intersection of results

""" % (NAME, VERSION, NAME))

def esc(s):
	return s.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')

def unique(s):
	u = {}
	try:
		for x in s:
			u[x] = 1
	except TypeError:
		del u  # move on to the next method
	else:
		return u.keys()
	return []

def intersection(s1, s2):
#	if len(s1) == 0 and len(s2) > 0:
#		return s2
#	elif len(s1) > 0 and len(s2) == 0:
#		return s1
#	else:
		return filter(lambda x:x in s1,s2)

def get_album_info(album_ids):
	query = "WHERE id IN (%s)" % ",".join(album_ids)
	cmd = """dcop amarok collection query "SELECT id,';',name,'|' from album %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	for line in lines:
		words = string.split(line.rstrip("\n"), ";")
		albumnames[words[0]] = words[1]

def get_artist_info():
	cmd = """dcop amarok collection query "SELECT id,';',name,'|' from artist"|tr -d "\n"|sed -e 's/|/\\n/g' """
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	for line in lines:
		words = string.split(line.rstrip("\n"), ";")
		artistnames[words[0]] = words[1]

def get_genre_info():
	cmd = """dcop amarok collection query "SELECT id,';',name,'|' from genre"|tr -d "\n"|sed -e 's/|/\\n/g' """
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	for line in lines:
		words = string.split(line.rstrip("\n"), ";")
		genrenames[words[0]] = words[1]

def get_year_info():
	cmd = """dcop amarok collection query "SELECT id,';',name,'|' from year"|tr -d "\n"|sed -e 's/|/\\n/g' """
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	for line in lines:
		words = string.split(line.rstrip("\n"), ";")
		yearnames[words[0]] = words[1]


def get_tracks(album_ids):
	albumdict = {}

	query = "WHERE album in (%s)" % ",".join(album_ids)
	cmd = """dcop amarok collection query "SELECT album,';',title,';',artist,';',track,';',length,';',year,';',genre,'|' FROM tags %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	for line in lines:
		words = string.split(line.rstrip("\n"), ";")
		id = words[0]
		title = words[1]
		artist = words[2]
		if artistnames.has_key(artist):
			artist = artistnames[artist]
		if len(words[3]) == 0:
			continue
		track = "%03d" % int(words[3])
		length = words[4]
		year = words[5]
		genre = words[6]
		info = { "title" : title, "artist" : artist, "length" : length }
		if albumdict.has_key(id):
			albumdict[id]["tracks"][track] = info
		else:
			albumdict[id] = {"tracks" : {track : info} }

		if albumdict[id].has_key("artist"):
			if albumdict[id]["artist"] != artist:
				albumdict[id]["artist"] = "(Various)"
		else:
			albumdict[id]["artist"] = artist

		if albumdict[id].has_key("year"):
			if albumdict[id]["year"] != year:
				albumdict[id]["year"] = ""
		else:
			albumdict[id]["year"] = year

		if albumdict[id].has_key("genre"):
			if albumdict[id]["genre"] != genre:
				albumdict[id]["genre"] = ""
		else:
			albumdict[id]["genre"] = genre

	return albumdict

def get_artist_ids():
	query = "WHERE name LIKE '%%%s%%'" % options["artist"]
	cmd = """dcop amarok collection query "SELECT id,'|' FROM artist %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	return [line.rstrip("\n") for line in lines]

def get_genre_ids():
	query = "WHERE name LIKE '%%%s%%'" % options["genre"]
	cmd = """dcop amarok collection query "SELECT id,'|' FROM genre %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	return [line.rstrip("\n") for line in lines]

def search_albumtitles():
	query = "WHERE name LIKE '%%%s%%'" % options["albumtitle"]
	cmd = """dcop amarok collection query "SELECT id,'|' FROM album %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	return [line.rstrip("\n") for line in lines]

def songtitle_query():
	if not options.has_key('songtitle'):
		return ""
	return " AND title LIKE '%%%s%%'" % options["songtitle"]

def artist_query():
	if not options.has_key('artist'):
		return ""
	artist_ids = get_artist_ids()
	if len(artist_ids) == 0:
		return ""
	return " AND artist IN (%s)" % ",".join(artist_ids)

def genre_query():
	if not options.has_key('genre'):
		return ""
	genre_ids = get_genre_ids()
	if len(genre_ids) == 0:
		return ""
	return " AND genre IN (%s)" % ",".join(genre_ids)

def search_tags():
	query = ""
	query += artist_query()
	query += genre_query()
	# I can't figure out how to make the songtitle an OR query right now
	# so assume that if -t is specified, along with another option, we're only
	# searching album titles
	if len(query) == 0:
		query += songtitle_query()
	query = "WHERE 1" + query
	cmd = """dcop amarok collection query "SELECT album,'|' FROM tags %s"|tr -d "\n"|sed -e 's/|/\\n/g' """ % query
	if options.has_key('verbose'):
		sys.stderr.write(cmd)
	lines = os.popen(cmd).readlines()
	return [line.rstrip("\n") for line in lines]

def get_albums():
	album_ids = []
	if options.has_key("albumtitle"):
#		album_ids[:0] = search_albumtitles()
		album_ids = search_albumtitles()
	if options.has_key("songtitle") or options.has_key("artist") or options.has_key("genre"):
		tmp_ids = search_tags()
		if len(album_ids) > 0:
			album_ids = intersection(album_ids, tmp_ids)
		else:
			album_ids = tmp_ids
	get_album_info(album_ids)
	get_artist_info()
	get_genre_info()
	get_year_info()
	albumdict = get_tracks(album_ids)
	return albumdict

def print_entry(id, album):
	year = album["year"]
	if len(year) > 0:
		year = esc(yearnames[year])
	genre = album["genre"]
	if len(genre) > 0:
		genre = esc(genrenames[genre])
	out = "<entry id='%s'>\n<title>%s</title>\n<artist>%s</artist>\n<year>%s</year>\n<genre>%s</genre>\n" \
	  % (id, esc(albumnames[id]), esc(album["artist"]), year, genre)

	if not options.has_key("noimages"):
		imagename = md5.new( "%s%s" % (album["artist"].lower(), albumnames[id].lower()) ).hexdigest()
		if len(imagename) > 0 and os.path.isfile(coverdir + imagename):
			out += "<cover>%s.PNG</cover>\n" % imagename
			images.add(imagename)

	out += "\t<tracks>\n"
	tracks = album["tracks"].keys()
	tracks.sort()
	currTrack = 1
	for num in tracks:
		while currTrack < int(num):
			out += "<track></track>\n"
			currTrack += 1
		info = album["tracks"][num]
		mins,secs = divmod(int(info["length"]), 60)
		out += "\t\t<track>%s::%s::%s:%02d</track>\n" % (esc(info["title"]), esc(info["artist"]), mins, secs)
		currTrack = int(num) + 1
	out += "\t</tracks>\n"
	out += "</entry>\n"
	return out

def generate_images():
	out = "<images>\n"
	for img in images:
		filename = coverdir + img
		try:
			file = open(filename, "rb")
			data = file.read();
			out += "<image format='PNG' id='%s'>\n" % (img + '.PNG')
			out += base64.encodestring(data)
			out += '</image>\n'
		except IOError:
			pass
	out += "</images>\n"
	return out

def generate_xml():
	out = """<?xml version='1.0' encoding='UTF-8'?>
<!DOCTYPE tellico PUBLIC '-//Robby Stephenson/DTD Tellico V9.0//EN'
	                     'http://periapsis.org/tellico/dtd/v9/tellico.dtd'>
<tellico xmlns='http://periapsis.org/tellico/' syntaxVersion='9'>
	<collection title='amaroK' type='4'>
		<fields>
			<field name='_default'/>
		</fields>
"""

	albums = get_albums()
	for id in albums.keys():
		out += print_entry(id, albums[id])

	if not options.has_key("noimages"):
		out += generate_images()
	out += "</collection></tellico>\n"
	return out


if __name__ == "__main__":

	try:
		opts, args = getopt.getopt(sys.argv[1:], 'k:a:t:xqvh')
	except getopt.error, msg:
		print msg
		usage()
		sys.exit(2)

	dosearch = False
	for opt, arg in opts:
		if opt == '-h':
			usage()
			sys.exit(0)
		elif opt == '-v':
			options['verbose'] = 1
		elif opt == '-q':
			options['quiet'] = 1
		elif opt == '-x':
			options['noimages'] = 1
		# now if there's no argument, skip it
		elif len(arg) == 0:
			continue
		elif opt == '-t':
			options['albumtitle'] = arg
			options['songtitle'] = arg
			dosearch = True
		elif opt == '-a':
			options['artist'] = arg
			dosearch = True
		elif opt == '-k':
			options['genre'] = arg
			dosearch = True

	if not dosearch:
		usage()
		sys.exit(2)

	try:
		coverdir = os.environ['KDEHOME'];
	except KeyError:
		coverdir = os.environ['HOME'] + '/.kde';
	coverdir += "/share/apps/amarok/albumcovers/large/"
	xml = generate_xml()
	lang, enc = locale.getdefaultlocale()
	if not options.has_key('quiet'):
		print xml.decode(enc).encode('utf-8')

