symuid/symuid-import-itunes

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import datetime as dt
import dateutil.parser
import mutagen
import mutagen.id3
import mutagen.mp4
import os
import symuid.library.itunes
import urllib.parse
import xml.etree.ElementTree


def generate_play_count_tag_label(player, library_id, reg_dt):
    return 'symuid:pcnt:{}:{}:{}'.format(player, library_id, int(reg_dt.timestamp()))


def get_itunes_dict_value_node(dict_node, key):
    assert isinstance(dict_node, xml.etree.ElementTree.Element)
    assert isinstance(key, str)
    # WORKAROUND method getnext() is sadly not available
    for child_idx, child_node in enumerate(dict_node):
        if child_node.tag == 'key' and child_node.text == key:
            return dict_node[child_idx + 1]
    raise KeyError()


def get_itunes_dict_value(dict_node, key):
    value_node = get_itunes_dict_value_node(dict_node, key)
    if value_node.tag == 'string':
        return value_node.text
    elif value_node.tag == 'integer':
        return int(value_node.text)
    elif value_node.tag == 'date':
        return dateutil.parser.parse(value_node.text)
    else:
        return value_node


def set_play_count_tag(track_path, player, library_id, reg_dt, play_count):
    assert isinstance(reg_dt, dt.datetime), reg_dt
    tag_label = generate_play_count_tag_label(
        player=player,
        library_id=library_id,
        reg_dt=reg_dt,
    )
    track = mutagen.File(filename=track_path)
    if isinstance(track.tags, mutagen.id3.ID3):
        tag_label_id3 = 'TXXX:' + tag_label
        if not tag_label_id3 in track.tags:
            # mutagen.id3._specs.EncodedTextSpec.write encodes
            # 'desc' and 'text'
            tag = mutagen.id3.TXXX(
                encoding=mutagen.id3.Encoding.LATIN1,
                desc=tag_label,
                text=[str(play_count)],
            )
            track.tags.add(tag)
            track.save()
            print('{!r}: set ID3 tag {!r}'.format(track_path, tag))
    elif isinstance(track.tags, mutagen.mp4.MP4Tags):
        tag_label_mp4 = '----:' + tag_label
        if not tag_label_mp4 in track.tags:
            track.tags[tag_label_mp4] = tag = mutagen.mp4.MP4FreeForm(
                # "a signed big-endian integer with length one of { 1,2,3,4,8 } bytes"
                # TODO set byte length properly
                data=play_count.to_bytes(1, byteorder='big'),
                dataformat=mutagen.mp4.AtomDataType.INTEGER,
            )
            track.save()
            print('{!r}: set MP4 tag {!r}'.format(track_path, tag))
    else:
        raise Exception(track_path)


def symuid_import_itunes(xml_library_path, root_url, root_path):
    root_path = os.path.expanduser(root_path)
    lib = symuid.library.itunes.XmlLibrary(xml_library_path)
    for track_node in lib._root_dict['Tracks'].iterfind('./dict'):
        try:
            track_url = get_itunes_dict_value(track_node, 'Location')
        except KeyError:
            track_url = None
        try:
            play_count = get_itunes_dict_value(track_node, 'Play Count')
        except KeyError:
            play_count = 0
        try:
            last_play_dt = get_itunes_dict_value(track_node, 'Play Date UTC')
        except KeyError:
            last_play_dt = None
        # TODO create tag if last_play_dt is None
        if last_play_dt and track_url and track_url.startswith(root_url):
            track_path = os.path.join(
                root_path,
                urllib.parse.unquote(track_url[len(root_url):]),
            )
            # TODO dt=dt.datetime.now()
            set_play_count_tag(
                track_path=track_path,
                player='itunes',
                library_id=lib.id,
                reg_dt=last_play_dt,
                play_count=play_count,
            )


def _init_argparser():
    import argparse
    argparser = argparse.ArgumentParser(description=None)
    argparser.add_argument('xml_library_path')
    argparser.add_argument(
        '--root-url',
        default='file://localhost/',
        help='(default: %(default)r)',
    )
    argparser.add_argument(
        '--root-path',
        default='',
        help='(default: %(default)r)',
    )
    return argparser


def main(argv):
    argparser = _init_argparser()
    args = argparser.parse_args(argv[1:])
    symuid_import_itunes(**vars(args))
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main(sys.argv))