#!/usr/bin/python

# This script creates an index of all iTunes tracks and their respective lyrics, and provides a fast regex-based search function.
# Requirements: Mac OS X 10.5 or higher, iTunes.

# INSTALLATION:
# 1) Copy search.py whereever you want (make sure that your user has write permission, though)
# 2) (optional) Create a link in /usr/bin or some other directory in your PATH
# 3) Create an index of your lyrics
# 4) Start using it!
# e.g.
# 1) mkdir -p ~/Applications/Lyricsearch && cp search.py ~/Applications/Lyricsearch
# 2) (as root) ln -s ~/Applications/Lyricsearch/search.py /usr/bin/lyricsearch
# 3) Do "lyricsearch --index"
# 4) Search by doing "lyricsearch 'search string'"

# Copyright (c) Thomas Backman <serenity@exscape.org>, 2009. Written 2009-06-16 (evening) - 2009-06-17 (morning).
# Released under a new BSD license. Do pretty much what you want, but credit me.

# --display capability added in about 15 minutes, 2010-01-20.

import sys, pickle, re, getopt
import os.path

#sys.setdefaultencoding('utf-8')

# Paths to the script and index file
dirpath = os.path.dirname(os.path.realpath(__file__))
indexpath = os.path.join(dirpath, "lyric_index.dat")

def usage():
	print """usage: script --index, or
script [--verbose / -v] [--word / -w] 'regex to search for' | [--display / -d id ]"""
	sys.exit(1)

try:
	opts, args = getopt.gnu_getopt(sys.argv[1:], "divw", ["display", "index", "verbose", "word"])
except getopt.GetoutError, err:
	print str(err)
	usage()

index = False
verbose = False
word = False
arg = ""
display = False

for o, a in opts:
	if o in ("-v", "--verbose"):
		verbose = True
	elif o in ("-i", "--index"):
		index = True
	elif o in ("-h", "--help"):
		usage()
	elif o in ("-w", "--word"):
		word = True
	elif o in ("-d", "--display"):
		display = True
	else:
		usage()

if index and len(args) != 0:
	usage() # We can't index AND search
elif len(args) != 1 and not index:
	usage() # We can't search without something to search for
elif not index:
	arg = args[0]

if index:
	from Foundation import *
	from ScriptingBridge import *

	iTunes = SBApplication.applicationWithBundleIdentifier_("com.apple.iTunes")

	if iTunes == None:
		print 'Unable to get a handle to iTunes. Make sure that iTunes is running (and that you are running Mac OS X 10.5 or higher)!'
		sys.exit(1)

	try:
		trackiterator = iTunes.sources()[0].playlists()[0].tracks()
	except:
		print 'Unable to get list of tracks'
		sys.exit(1)
	if trackiterator == None:
		print 'Unable to get list of tracks'
		sys.exit(1)

	num_tracks = len(trackiterator)
	all_tracks = []
	i = 0
	for track in trackiterator:
		t = track
		lyr = t.lyrics()
		if lyr == None:
			continue
		name = t.name()
		artist = t.artist()
		current_track = {'id': i, 'name': name, 'artist': artist, 'genre': t.genre(), 'lyrics': lyr.replace('\r','\n')}
		all_tracks.append(current_track)
		if verbose:
			print '>> Indexed [#%d] %s - %s' % (i, artist, name)
		i += 1
		if i % 100 == 0:
			print i, 'out of', num_tracks, 'tracks indexed (%.2f%%)' % ((float(i)/num_tracks)*100)

	try:
		f = open(indexpath, 'wb')
		pickle.dump(all_tracks, f, 2)
		f.close()
		print 'Done!'
		sys.exit(0)
	except Exception, e:
		print 'Unable to open file and write index! Error:'
		print str(e)
		sys.exit(1)

## Common to dispay and search ##
try:
	all_tracks = pickle.load(open(indexpath, 'rb'))
except:
	print 'Unable to read indexed data from lyric_index.dat! Have you run the script with the --index option?'
	sys.exit(1)

######################
# End of the index part, start of the display part
######################

if display:
	for track in all_tracks:
		if track['id'] == int(arg):
			print "Displaying lyrics for track #%d (%s - %s)\n" % (track['id'], track['artist'].encode('utf-8', 'ignore'), track['name'].encode('utf-8', 'ignore'))
			print track['lyrics'].encode('utf-8', 'ignore')
			sys.exit(0)
	print 'No track matching that ID was found!'
	sys.exit(1)
	

######################
# End of the diplay part, start of the search part
######################

try:
	if word:
		arg = r'\b' + arg + r'\b'
	regex = re.compile(arg, re.IGNORECASE | re.DOTALL | re.MULTILINE | re.LOCALE)
except:
	print 'Unable to compile regular expression! Make sure that your search string is a valid regular expression.'
	usage()

first = True # To make empty lines between songs, but not before the very first one
for track in all_tracks:
	lyr = track['lyrics']
	if lyr == None:
		continue
	lyr = lyr.encode('utf-8', 'ignore')
	m = regex.search(lyr)
	if m:
		id = track['id']
		title = track['name'].encode('utf-8', 'ignore')
		artist = track['artist'].encode('utf-8', 'ignore')
		if first == False and verbose:
			print '' # Empty line between songs
		print "[#%d] %s - %s" % (id, artist, title)
		first = False

		if verbose:
			for line in lyr.split('\n'):
				if line.find(m.group(0)) != -1:
					print '  >> "%s"' % (line)