#!/usr/bin/python

import sys
import getopt
try:
    import musicbrainz
except ImportError, e:
    print >>sys.stderr, 'import: %s' % e
    sys.exit(1)

import lastfm.mbutil

class Abby(musicbrainz.mb):
    def __init__(self):
        musicbrainz.mb.__init__(self)

    def readdisc(self, dev='/dev/cdrom'):
        self.SetDevice(dev)
        try:
            self.Query(musicbrainz.MBQ_GetCDTOC)
            return self.GetResultData(musicbrainz.MBE_TOCGetCDIndexId)
        except musicbrainz.MusicBrainzError:
            print >>sys.stderr, ("cdrom: can't read TOC from disc")
            sys.exit(1)

    def query(self, discid=None, albumid=None):
        try:
            if discid:
                self.SetDepth(2)
                self.QueryWithArgs(musicbrainz.MBQ_GetCDInfoFromCDIndexId,
                    [discid])
            else:
                self.SetDepth(4)
                self.QueryWithArgs(musicbrainz.MBQ_GetAlbumById, [albumid])
            return self.querieddiscinfo()
        except musicbrainz.MusicBrainzError:
            return self.localdiscinfo()

    def querieddiscinfo(self):
        matches = self.GetResultInt(musicbrainz.MBE_GetNumAlbums)
        if matches == 0: return self.localdiscinfo()

        self.Select1(musicbrainz.MBS_SelectAlbum, 1)
        n = self.GetResultInt(musicbrainz.MBE_AlbumGetNumTracks)
        album = self.GetResultData(musicbrainz.MBE_AlbumGetAlbumName)

        artistid = self.GetIDFromURL(self.GetResultData(
            musicbrainz.MBE_AlbumGetAlbumArtistId))
        if artistid == musicbrainz.MBI_VARIOUS_ARTIST_ID:
            albumartist = 'Various Artists'
        else:
            # XXX: This is wrong! SG5DR.
            albumartist = self.GetResultData1(
                musicbrainz.MBE_AlbumGetArtistName, 1)

        tracks = []
        for i in range(1, n+1):
            tracks.append({
                'number': i,
                'album': album,
                'artist': self.GetResultData1(
                    musicbrainz.MBE_AlbumGetArtistName, i),
                'title': self.GetResultData1(
                    musicbrainz.MBE_AlbumGetTrackName, i),
                'length': self.GetResultInt1(
                    musicbrainz.MBE_AlbumGetTrackDuration, i) / 1000,
                'mbid': self.GetIDFromURL(self.GetResultData1(
                    musicbrainz.MBE_AlbumGetTrackId, i)),
                })

        return album, albumartist, tracks

    def localdiscinfo(self):
        # Ensure that previous query was local TOC (ugh)
        self.readdisc()

        first = self.GetResultInt(musicbrainz.MBE_TOCGetFirstTrack)
        last = self.GetResultInt(musicbrainz.MBE_TOCGetLastTrack)

        tracks = []
        for i in range(first, last+1):
            tracks.append({
                'number': i,
                'artist': '',
                'album': '',
                'title': '',
                'length': self.GetResultInt1(
                    musicbrainz.MBE_TOCGetTrackNumSectors, i+1) / 75,
                })

        return None, None, tracks

def freedb_dump(albumartist, album, discid, tracks):
    print '# fake CD database file generated by mbget'
    print '#'
    print '# Track frame offsets:'

    # Assume standard pregap
    total_len = 2
    for t in tracks:
        print '#       %d' % (total_len * 75)
        total_len += t['length']

    if albumartist and album:
        title = '%s / %s' % (albumartist, album)
    else:
        title = ''

    print '#'
    print '# Disc length: %d seconds' % total_len
    print '#'
    print '# Revision: 0'
    print '# Processed by: MusicBrainz'
    print '# Submitted by: MusicBrainz'
    print 'DISCID=%s' % discid
    print 'DTITLE=%s' % title

    for i, t in enumerate(tracks):
        if albumartist:
            if albumartist == t['artist']:
                title = t['title']
            else:
                title = '%s / %s' % (t['artist'], t['title'])
        else:
            title = ''
        print 'TTITLE%d=%s' % (i, title)

    print 'EXTD='

    for i, t in enumerate(tracks):
        print 'EXTD%d=' % i

    print 'PLAYORDER='
    print '.'

if __name__ == '__main__':
    shortopts = 'd:pi:a:ft'
    longopts = ['device=', 'print-discid', 'discid=', 'albumid=', 'freedb', 'template']

    try:
        opts, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
    except getopt.GetoptError, e:
        print >>sys.stderr, 'getopt: %s' % e
        sys.exit(1)

    sys.stdout = lastfm.mbutil.EncodedFile(sys.stdout)

    abby = Abby()
    discid = None
    albumid = None
    freedb = False
    template = False

    for opt, arg in opts:
        if opt in ('--device', '-d'):
            discid = abby.readdisc(arg)
        elif opt in ('--print-discid', '-p'):
            print abby.readdisc(arg)
            sys.exit(0)
        elif opt in ('--discid', '-i'):
            discid = arg
        elif opt in ('--albumid', '-a'):
            albumid = arg
        elif opt in ('--freedb', '-f'):
            freedb = True
        elif opt in ('--template', '-t'):
            template = True

    if not (discid or albumid): discid = abby.readdisc()
    album, albumartist, tracks = abby.query(discid, albumid)

    if (album and albumartist) or template:
        if freedb:
            freedb_dump(albumartist, album, discid, tracks)
        else:
            import lastfm.marshaller
            if args:
                tracks = [tracks[int(a)-1] for a in args]
            for t in tracks:
                # The serialized format is UTF-8, but when printing, we need
                # to encode to stdout instead... this is ugly.
                print lastfm.marshaller.dump(t).decode('utf-8')
    else:
        sys.exit(1)
