#!/usr/bin/env python
# encoding: utf-8
"""
update_tags.py

Created by Anchal Nigam on 2012-05-03.
Copyright (c) 2012 __MyCompanyName__. All rights reserved.
"""
# http://cl.ly/C3kK


import os, sys
import subprocess
from struct import unpack
from time import strftime
import urllib
import json
from math import log10, ceil
import dateutil.parser
import re

# some config options
DATE_TIME_FORMAT = "[ %Y-%m-%d @ %H:%M:%S %Z ]  "
DATE_FORMAT = "%Y-%m-%d"
SEPARATOR = " : "

# number of files to check against iTunes
BATCH_SIZE = 100

# add iTunes comments
ADD_OLD_TAGS_TO_COMMENTS = True

PLAYLIST_NO_ITUNES_TRACK_ID = "Z1: No iTunes Info"
PLAYLIST_DUPLICATE_TRACKS = "Z2: Duplicate Songs"
PLAYLIST_WITH_APPLE_INFO = "Z3: Updated Songs"
PLAYLIST_WITHOUT_APPLE_INFO = "Z4: No iTunes Information Found"

# do we want to show debug
SHOW_DEBUG_OUTPUT = True
SHOW_DEBUG_URL = True

# do we want to validate the countries instead
VALIDATE_COUNTRIES = False

# the url we use to get information from Apple
APPLE_LOOKUP_URL = "http://itunes.apple.com/lookup?id=%(id)s&country=%(country)s"

TEMP_FILE = "app.tmp"

# an array of countries to check
COUNTRIES = "US GB AU FR DE CA IT JP DZ AO AI AG AR AM AT AZ BS BH BD BB BY BE BZ BM BO BW BR BN BG CM KY CL CN CO CR CI HR CY CZ DK DM DO EC EG SV EE ET FI GH GR GD GT GY HN HK HU IS IN ID IE IL JM JO KZ KE KR KW LV LB LY LI LT LU MO MK MG MY MV ML MT MU MX MD MS MM NP NL NZ NI NE NG NO OM PK PA PY PE PH PL PT QA RO RU KN LC VC SA SN RS SG SK SI ZA ES LK SR SE CH TW TZ TH TT TN TR TC UG UA AE UY UZ VE VN VG YE AA AB AC AD AF AH AJ AK AL AN AP AQ AS AV AW AX AY BA BC BF BI BJ BK BL BP BQ BT BU BV BX CB CC CD CE CF CG CJ CK CP CQ CS CT CU CV CW CX DA DB DC DD DF DG DH DI DJ DL DN DP DQ DR DS DT DU DV DW DX DY EA EB ED EF EH EI EJ EK EL EM EN EO EP EQ ER EU EV EW EX EY EZ FA FB FC FD FE FF FG FH FJ FK FL FM FN FO FP FQ FS FT FU FV FW FX FY FZ GA GC GE GF GG GI GJ GK GL GM GN GO GP GQ GS GU GV GW GX GZ HA HB HC HD HE HF HG HH HI HJ HL HM HO HP HQ HS HT HV HW HX HY HZ IA IB IC IF IG IH II IJ IK IM IO IP IQ IR IU IV IW IX IY IZ JA JB JC JD JE JF JG JH JI JJ JK JL JN JQ JR JS JT JU JV JW JX JY JZ KA KB KC KD KF KG KH KI KJ KK KL KM KO KP KQ KS KT KU KV KX LA LD LE LF LG LH LJ LL LM LN LO LP LQ LR LS LW LX LZ MA MB MC ME MF MH MI MJ MN MP MQ MR MW MZ NA NB NC ND NF NH NJ NK NM NN NQ NR NS NT NU NV NW NX NY OA OB OC OD OE OF OG OH OI OJ OK OL ON OO OP OQ OR OS OT OU OV OW OX OY OZ PB PC PD PF PG PI PJ PM PN PO PP PQ PR PS PU PV PW PX PZ QB QC QD QE QF QG QH QI QJ QK QL QM QN QO QP QQ QR QS QT QU QV QW QX QY QZ RA RB RC RD RE RF RG RH RI RJ RK RL RM RN RP RQ RR RT RV RW RX RY RZ SB SC SD SF SH SJ SL SM SO SP SQ SS ST SU SW SX SY SZ TA TB TD TE TF TG TI TJ TK TL TM TO TP TQ TS TU TV TX TY UB UC UD UE UF UH UI UJ UK UL UM UN UO UP UQ UR UT UU UV UW UX VA VB VD VF VH VI VJ VK VL VM VO VP VQ VR VS VT VU VV VW VX VY VZ WA WB WC WD WE WF WG WH WI WJ WK WL WM WN WO WP WQ WR WS WT WU WV WW WX WY WZ XA XB XC XD XE XF XG XH XI XJ XK XL XM XN XO XP XQ XR XS XT XU XV XW XX XY XZ YA YB YC YD YF YG YH YI YJ YK YL YM YN YO YP YQ YR YS YT YU YV YW YX YY YZ ZB ZC ZD ZE ZF ZG ZH ZI ZJ ZK ZL ZM ZN ZO ZP ZQ ZR ZS ZT ZU ZV ZW ZX ZY ZZ".split( " " )
COUNTRIES = "US GB AU FR DE CA IT JP DZ AO AI AG AR AM AT AZ BS BH BD BB BY BE BZ BM BO BW BR BN BG CM KY CL CN CO CR CI HR CY CZ DK DM DO EC EG SV EE ET FI GH GR GD GT GY HN HK HU IS IN ID IE IL JM JO KZ KE KR KW LV LB LY LI LT LU MO MK MG MY MV ML MT MU MX MD MS MM NP NL NZ NI NE NG NO OM PK PA PY PE PH PL PT QA RO RU KN LC VC SA SN RS SG SK SI ZA ES LK SR SE CH TW TZ TH TT TN TR TC UG UA AE UY UZ VE VN VG YE".split( " " )

# print log
def log( *messages, **options ):
    # if we're appending we don't need the date
    date = SEPARATOR if "append" in options and options["append"] else strftime( DATE_TIME_FORMAT )
    # if we have a new line
    if "newline" in options and not options["newline"]:
        sys.stdout.write( date + SEPARATOR.join( map( str, messages ) ) )
    # normal
    else:
        print date + SEPARATOR.join( map( str, messages ) )

# print debug info
def debug( *messages, **options ):
    # only if we want to show debug
    if SHOW_DEBUG_OUTPUT:
        # if we're appending we don't need to add DEBUG to start
        if "append" not in options or not options["append"]:
            messages = ( "DEBUG", ) + messages
        # call normal print
        log( *messages, **options )

# gets selected tracks and returns an array of data base id and path of each track
def getiTunesMatchedTracks():
    applescript = """
on replace_chars(this_text, search_string, replacement_string)
    set AppleScript\'s text item delimiters to the search_string
    set the item_list to every text item of this_text
    set AppleScript\'s text item delimiters to the replacement_string
    set this_text to the item_list as string
    set AppleScript\'s text item delimiters to ""
    return this_text
end replace_chars
tell application "iTunes"
    set output to "["
    repeat with t in selection
        if kind of t is "Matched AAC audio file" then
            set t_location to location of t
            set t_path to POSIX path of t_location
            set output to output & "{"
            set output to output & "\'database id\':" & (database ID of t) & ", "
            set output to output & "\'path\':\\"" & t_path & "\\""
            set output to output & "}, "
        end if
    end repeat
    set output to output & "]"
    get output
end tell
"""
    # opens a process to call applescript
    process = subprocess.Popen( [ "osascript", "-e", applescript ], stdout = subprocess.PIPE, stderr = subprocess.STDOUT )
    
    # gets output
    ( stdout, stderr ) = process.communicate()
    
    # assuming everything went good
    if( process.wait() == 0 ):
        # evalulate the returned data and return
        return eval( stdout.strip() )
    
    # this only happens if the applescript failed
    return []

# determine how much padding to use provided a max number
def paddingWidth( number ):
    return int( ceil( log10( number ) ) )

# get JSON data from a URL
def getJSONDataFromURL( url ):
    return json.loads( urllib.urlopen( url ).read() )

# validate the list of countries
def validateCountries():
    # make a copy of the countries
    countries_copy = COUNTRIES[:]
    
    number_of_countries = str( len( countries_copy ) )
    padding = paddingWidth( len( countries_copy ) )
    
    debug( "checking %s countries" % ( number_of_countries, ) )
    
    # iterate through each country
    for ( i, country ) in enumerate( countries_copy ):
        # get a URL for that country
        url = APPLE_LOOKUP_URL % { "id": 1, "country": country }
        
        debug( "checking country", str( i + 1 ).rjust( padding ) + " of " + number_of_countries, country, url )
        
        # get the JSON for that url and see if there was an error
        if "errorMessage" in getJSONDataFromURL( url ):
            # remove the country from the list if there was an error
            COUNTRIES.remove( country )
    
    # show the left over/valid countries
    debug( "found %s valid countries" % ( str( len( COUNTRIES ) ) ) )
    print " ".join( COUNTRIES )


def updateiTunesInformation( no_itunes_track_id, duplicate_tracks, with_apple_info, without_apple_info ):
    applescript = """
set no_itunes_track_id_playlist to "%s"
set duplicate_tracks_playlist to "%s"
set tracks_to_update_playlist to "%s"
set without_apple_info_playlist to "%s"
""" % ( PLAYLIST_NO_ITUNES_TRACK_ID, PLAYLIST_DUPLICATE_TRACKS, PLAYLIST_WITH_APPLE_INFO, PLAYLIST_WITHOUT_APPLE_INFO )
    log( "updating information in iTunes" )
    
    log( "adding " + str( len( no_itunes_track_id ) ) + " to playlist '" + PLAYLIST_NO_ITUNES_TRACK_ID + "'" )
    applescript += "set no_itunes_track_id to {"
    for database_track_id in no_itunes_track_id:
        applescript += '{id:%s},' % ( database_track_id, )
    if len( no_itunes_track_id ):
        applescript = applescript[:-1]
    applescript += "}\n"
    
    log( "adding " + str( len( duplicate_tracks ) ) + " to playlist '" + PLAYLIST_DUPLICATE_TRACKS + "'" )
    applescript += "set duplicate_tracks to {"
    for database_track_id in duplicate_tracks:
        applescript += '{id:%s},' % ( database_track_id, )
    if( len( duplicate_tracks ) ):
        applescript = applescript[:-1]
    applescript += "}\n"
    
    log( "updating tag information for " + str( len( with_apple_info ) ) + " tracks and adding them to playlist '" + PLAYLIST_WITH_APPLE_INFO + "'" )
    applescript += "set tracks_to_update to {"    
    for ( database_track_id, data ) in with_apple_info.items():
        url = data["lookupURL"]
        cleaned = data.get( "trackExplicitness", "" ) == "cleaned"
        artist = data.get( "artistName", "" ).replace( "\"", "\\\"" )
        album = data.get( "collectionName", "" ).replace( "\"", "\\\"" )
        track = data.get( "trackName", "" ).replace( "\"", "\\\"" )
        year = str( dateutil.parser.parse( data["releaseDate"] ).year ) if "releaseDate" in data else ""
        discCount = data.get( "discCount", "" )
        discNumber = data.get( "discNumber", "" )
        trackCount = data.get( "trackCount", "" )
        trackNumber = data.get( "trackNumber", "" )
        primaryGenraName = data.get( "primaryGenreName", "" ).replace( "\"", "\\\"" )
        
        if cleaned:
            track = track + " (cleaned)"
        
        a = '{id:%i, nArtist:"%s", nAlbum:"%s", nTrack:"%s", nYear:"%s", nDiscCount:"%s", nDiscNumber:"%s", nTrackCount:"%s", nTrackNumber:"%s", nGenre:"%s", lookupURL:"%s"},' % ( database_track_id, artist, album, track, year, discCount, discNumber, trackCount, trackNumber, primaryGenraName, url )
        
        applescript += '{id:%i, nArtist:"%s", nAlbum:"%s", nTrack:"%s", nYear:"%s", nDiscCount:"%s", nDiscNumber:"%s", nTrackCount:"%s", nTrackNumber:"%s", nGenre:"%s", lookupURL:"%s"},' % ( database_track_id, artist, album, track, year, discCount, discNumber, trackCount, trackNumber, primaryGenraName, url )
    
    if( len( with_apple_info ) ):
        applescript = applescript[:-1]
    applescript += "}\n"
    
    log( "adding " + str( len( without_apple_info ) ) + " to playlist '" + PLAYLIST_WITHOUT_APPLE_INFO + "'" )
    applescript += "set without_apple_info to {"
    for ( itunes_id, database_track_id ) in without_apple_info.items():
        url = APPLE_LOOKUP_URL % { "id" : itunes_id, "country": "" }
        applescript += '{id:%s,iTunesID:%s,lookupURL:"%s"},' % ( database_track_id, itunes_id, url )
    if( len( without_apple_info ) ):
        applescript = applescript[:-1]
    applescript += "}\n"
    
    
    applescript += """


tell application "iTunes"
    if not (exists (some user playlist whose name is no_itunes_track_id_playlist)) then make user playlist with properties {name:no_itunes_track_id_playlist}
    if not (exists (some user playlist whose name is duplicate_tracks_playlist)) then make user playlist with properties {name:duplicate_tracks_playlist}
    if not (exists (some user playlist whose name is tracks_to_update_playlist)) then make user playlist with properties {name:tracks_to_update_playlist}
    if not (exists (some user playlist whose name is without_apple_info_playlist)) then make user playlist with properties {name:without_apple_info_playlist}
    set no_itunes_track_id_playlist to playlist no_itunes_track_id_playlist
    set duplicate_tracks_playlist to playlist duplicate_tracks_playlist
    set tracks_to_update_playlist to playlist tracks_to_update_playlist
    set without_apple_info_playlist to playlist without_apple_info_playlist
    
    repeat with each_track in no_itunes_track_id
        set track_pointer to (item 1 of (every track of library playlist 1 whose database ID is (id of each_track)))
        add (get location of track_pointer) to no_itunes_track_id_playlist
    end repeat
    repeat with each_track in tracks_to_update
        set track_pointer to (item 1 of (every track of library playlist 1 whose database ID is (id of each_track)))"""
    
    if ADD_OLD_TAGS_TO_COMMENTS:
        applescript += """
        set comment of track_pointer to "Artist: " & (artist of track_pointer) & "
Album: " & (album of track_pointer) & "
Name: " & (name of track_pointer) & "
Year: " & (year of track_pointer) & "
Disc Count: " & (disc count of track_pointer) & "
Disc Number: " & (disc number of track_pointer) & "
Track Count: " & (track count of track_pointer) & "
Track Number: " & (track number of track_pointer) & "
Genre: " & (genre of track_pointer)"""
    
    applescript += """
        set artist of track_pointer to nartist of each_track
        set album of track_pointer to nalbum of each_track
        set name of track_pointer to ntrack of each_track
        set year of track_pointer to nyear of each_track
        set disc count of track_pointer to ndiscCount of each_track
        set disc number of track_pointer to ndiscNumber of each_track
        set track count of track_pointer to ntrackCount of each_track
        set track number of track_pointer to ntrackNumber of each_track
        set genre of track_pointer to nGenre of each_track
        set comment of track_pointer to (comment of track_pointer) & "
lookup URL: " & lookupURL of each_track
        add (get location of track_pointer) to tracks_to_update_playlist
    end repeat
    repeat with each_track in without_apple_info
        set track_pointer to (item 1 of (every track of library playlist 1 whose database ID is (id of each_track)))
        set comment of track_pointer to (comment of track_pointer) & "
{ 'oldiTunesID': " & (iTunesID of each_track) & ", 'lookUpURL': '" & (lookupURL of each_track) & "' }"
        add (get location of track_pointer) to without_apple_info_playlist
    end repeat
    repeat with each_track in duplicate_tracks
        set track_pointer to (item 1 of (every track of library playlist 1 whose database ID is (id of each_track)))
        add (get location of track_pointer) to duplicate_tracks_playlist
        set comment of track_pointer to (comment of track_pointer) & "
duplicate"
    end repeat
end tell


"""
    
    f = open( TEMP_FILE, "w" )
    f.write( applescript.encode('utf-8') )
    f.close()
    
    # opens a process to call applescript
    process = subprocess.Popen( [ "osascript", TEMP_FILE ], stdout = subprocess.PIPE, stderr = subprocess.STDOUT )
    
    # gets output
    ( stdout, stderr ) = process.communicate()
    
    # assuming everything went good
    if( process.wait() == 0 ):
        # evalulate the returned data and return
        return True
    
    # this only happens if the applescript failed
    return False
    
# main function
def main():
    log( "getting a list of selected iTunes matched tracks" )
    
    # get a list of selected tracks
    tracks = getiTunesMatchedTracks()
    
    # make sure we got some tracks
    if( len( tracks ) == 0 ):
        log( "no tracks selected", "exiting" )
        return
    
    # tracks that do not have an iTunes ID in the file data
    tracks_with_no_itunes_id = []
    
    # tracks with iTunes ID data
    tracks_with_itunes_id = {}
    
    # duplicates that have matching iTunes ID data
    duplicate_tracks = []
    
    number_of_tracks = str( len( tracks ) )
    padding = paddingWidth( len( tracks ) )
    
    log( "found " + str( len( tracks ) ), "looking for iTunes track ID" )
    
    # iterate through each track
    for ( i, track ) in enumerate( tracks ):
        track_database_id = track["database id"]
        track_path = track["path"]
        
        debug( "looking for iTunes track ID", str( i + 1 ).rjust( padding ) + " of " + number_of_tracks, track_path.replace( os.path.dirname( os.path.dirname( os.path.dirname( track_path ) ) ), "" ), newline = False )
        
        # look for the iTunes track ID which should be in the first 1024 bytes
        data = open( track_path, "rb" ).read( 1024 )
        index = data.index( "song" )
        
        # we found data
        if( index ):
            # unpack the iTunes track ID
            track_itunes_id = unpack( "!I", data[index+4:index+4+4] )[0]
            debug( "found iTunes track ID: " + str( track_itunes_id ), append = True )
            
            # if it already in the list that means we have a duplicate
            if( track_itunes_id in tracks_with_itunes_id ):
                # add both the new one and the duplicate to the duplicate list
                duplicate_tracks.append( track_database_id )
                if tracks_with_itunes_id[track_itunes_id] not in duplicate_tracks:
                    duplicate_tracks.append( tracks_with_itunes_id[track_itunes_id] )
            else:
                # add it to the list to be processed later
                tracks_with_itunes_id[track_itunes_id] = track_database_id
        else:
            # we couldn't find an iTunes track ID
            debug( "no iTunes track ID found", append = True )
            tracks_with_no_itunes_id.append( track_database_id )
    
    log( "found " + str( len( tracks_with_no_itunes_id ) ) + " tracks with no iTunes track ID" )
    log( "found " + str( len( duplicate_tracks ) ) + " duplicate tracks" )
    log( "found " + str( len( tracks_with_itunes_id ) ) + " tracks with iTunes track ID", "checking for information from Apple in batches of " + str( BATCH_SIZE ) )
    
    tracks_with_apple_info = {}
    tracks_without_apple_info = []
    all_itunes_ids = tracks_with_itunes_id.keys()
    
    num_countries = str( len( COUNTRIES ) )
    
    padding_countries = paddingWidth( len( COUNTRIES ) )
    padding_batch = paddingWidth( int( ceil( float( len( all_itunes_ids ) ) / BATCH_SIZE ) ) )
    
    # check batches of files in every country
    for ( i, country ) in enumerate( COUNTRIES ):
        all_itunes_ids_copy = all_itunes_ids[:]
        if( not len( all_itunes_ids_copy ) ):
            continue
        #log( "checking country", country + " (" + str( i + 1 ).rjust( padding_countries ) + " of " +  num_countries + ")" )
        
        # make a copy so we can delete/remove matched files
        num_tracks = str( len( all_itunes_ids_copy ) )
        num_batches = str( int( ceil( float( len( all_itunes_ids_copy ) ) / BATCH_SIZE ) ) )
        
        start = 0
        batch = 1
        
        # go through each batch
        while start < len( all_itunes_ids_copy ):
            end = start + BATCH_SIZE
            
            # get a URL
            url = APPLE_LOOKUP_URL % { "id": ",".join( map( str, all_itunes_ids_copy[start:end] ) ), "country": country }
            
            debug( "checking country", country + " (" + str( i + 1 ).rjust( padding_countries ) + " of " + num_countries + ")", "batch " + str( batch ).rjust( padding_batch ) + " of " + num_batches, url if SHOW_DEBUG_URL else None, newline = False )
            
            # get the JSON from Apple
            data = getJSONDataFromURL( url )
            debug( "found " + str( data["resultCount"] ) + " tracks", append = True )
            
            # did we find data
            if data["resultCount"]:
                # iterate through each found data
                for ( j, result ) in enumerate( data["results"] ):
                    # get the track ID from the data
                    track_itunes_id = result["trackId"]
                    # get the database ID
                    track_database_id = tracks_with_itunes_id[track_itunes_id]
                    # remove it from the list cause we found it
                    all_itunes_ids.remove( track_itunes_id )
                    # add lookup URL info
                    result["lookupURL"] = APPLE_LOOKUP_URL % { "id" : result["trackId"], "country": country }
                    # save it to the list of found ones
                    tracks_with_apple_info[track_database_id] = result
                    # remove it from the master list
                    del tracks_with_itunes_id[track_itunes_id]
            batch = batch + 1
            start = end
        
        num_tracks_found = len( all_itunes_ids_copy ) - len( all_itunes_ids )
        num_tracks_not_found = len( all_itunes_ids )
        
        log( "checking country", country + " (" + str( i + 1 ).rjust( padding_countries ) + " of " +  num_countries + ")", "tracks found ", str( num_tracks_found ) + " of " + str( len( all_itunes_ids_copy ) ), str( num_tracks_not_found ) + " left to go" )
    
    tracks_without_apple_info = tracks_with_itunes_id
    
    log( "found " + str( len( tracks_with_apple_info ) ) + " with information from Apple" )
    log( "there are " + str( len( tracks_without_apple_info ) ) + " with no information at Apple" )
    
    # tracks with no itunes track id   tracks_with_no_itunes_id
    # duplicate tracks           duplicate_tracks
    # tracks with info     tracks_with_apple_info
    # tracks without info    tracks_without_apple_info
    
    # update iTunes with the information we collected
    if updateiTunesInformation( no_itunes_track_id = tracks_with_no_itunes_id, duplicate_tracks = duplicate_tracks, with_apple_info = tracks_with_apple_info, without_apple_info = tracks_without_apple_info ):
        log( "finished with no issues" )
    else:
        log( "finished with errors" )
    
if __name__ == '__main__':
    if VALIDATE_COUNTRIES:
        validateCountries()
    else:
        main()