jmoney

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

import os
import sys
import six
import glob
import stat
import yaml
import shutil
import logging
import tempfile
from time import sleep
from argparse import ArgumentParser, ArgumentTypeError, ArgumentDefaultsHelpFormatter, FileType
from subprocess import call

try:
    import chardet
    HAS_CHARDET = True
except ImportError:
    HAS_CHARDET = False
try:
    import DiscID, CDDB
    HAS_DISCID = Ture
except ImportError:
    HAS_DISCID = False
try:
    import discogs_client
    HAS_DISCOGS = True
except ImportError:
    HAS_DISCOGS = False
try:
    shutil.get_terminal_size
    os.environ['COLUMNS'] = str(shutil.get_terminal_size().columns)
except AttributeError:
    pass  # Cannot set argparse output width to actual terminal size

get_input = raw_input if six.PY2 else input


logging.info('Begin')
log = logging.getLogger('jmoney')
exit_codes = {'disc_drive': 11,
              'freedb_query': 21,
              'disc_path': 31,
              'rip_disc': 41,
              'edit_disc': 51,
              'rename_tracks': 61}


def encode_dict(d, coding):
    '''
    Encode top level key and value strings in `d` in `encoding`
    '''
    transcoded_d = {}
    for k,v in d.items():
        try:
            if isinstance(k, six.string_types):
                k = k.decode(coding, 'replace').encode('UTF-8')
            if isinstance(v, six.string_types):
                v = v.decode(coding, 'replace').encode('UTF-8')
        except UnicodeDecodeError:
            pass
        transcoded_d[k] = v
    return transcoded_d


def suggest_encodings(entry):
    '''
    Try to use chardet to determine disc record encoding
    '''
    if HAS_CHARDET:
        # Get encoding of text fields
        detected = []
        text = b''
        for d in entry['disc_info'], entry['track_info']:
            for k, v in d.items():
                if isinstance(v, (six.string_types, bytes)):
                    text += v
        det = chardet.detect(text)
        return (det.get('encoding'), 'conf: {0:.2}'.format(det.get('confidence', 0)))
    else:
        return ['iso-8859-1', 'iso-8859-2', 'UTF-8']


class TextFormat:
    '''
    ANSI Select Graphic Rendition (SGR) code escape sequence formatting
    '''
    prefix = '\x1b['
    suffix = 'm'
    all_codes = {
        'reset': '0',

        'bold': '1',
        'faint': '2',
        'italic': '3',
        'underline': '4',

        'black': '30',
        'red': '31',
        'green': '32',
        'yellow': '33',
        'blue': '34',
        'magenta': '35',
        'cyan': '36',
        'white': '37',
        'extended': '38',
        'default': '39',
    }

    def __init__(self, *attrs, **kwargs):
        '''
        :param attrs: Are the attribute names of any format codes in
            ``all_codes``

        :param kwargs: May contain
            - ``reset``, prepend reset SGR code to sequence (default ``True``)

        Example:

        .. code-block:: python

            red_underlined = TextFormat('red', 'underline')
            print(
                '{0}Can you read this?{1}'
                ).format(red_underlined, TextFormat('reset'))
        '''
        self.codes = [self.all_codes[attr.lower()] for attr in attrs if isinstance(attr, six.string_types)]

        if kwargs.get('reset', True):
            self.codes[:0] = [self.all_codes['reset']]

        self.opts = kwargs.get('opts', {})

        if self.opts.get('no_color', False):
            self.sequence = ''
        else:
            self.sequence = '{}{}{}'.format(self.prefix, ';'.join(self.codes), self.suffix)

    def __call__(self, text, reset=True):
        '''
        Format :param text: by prefixing ``self.sequence`` and suffixing the
        reset sequence if :param reset: is ``True``.

        Example:

        .. code-block:: python

            green_text = TextFormat('blink', 'green')
            'The answer is: {0}'.format(green_text(42))
        '''
        end = TextFormat('reset') if reset else ''
        return '{}{}{}'.format(self.sequence, text, end)

    def __str__(self):
        return self.sequence

    def __repr__(self):
        return self.sequence


class Input:
    '''
    Prompt for and collect user input
    '''
    def __init__(self, opts):
        '''
        Setup opts
        '''
        self.opts = opts
        self.prompt_text = TextFormat('bold', 'green')
        self.invalid_text = TextFormat('bold', 'red')

    def __call__(self, prompt, input_type, valid=None, action=None):
        '''
        Prompt the user until valid input is received

        ``prompt``: The text presented to the user

        ``input_type``: A description of the kind of input expected

        ``valid``: An iterable of valid responses which are expected to be of
        type ``str`` or a single string, both in all lowercase

        ``action``: A function that is run with the collected input as an
        argument

        If ``action`` is given, the input and the results of ``action`` on the
        input are returned, otherwise only the input is returned
        '''
        while True:
            try:
                value = get_input(self.prompt_text(prompt))
                if action:
                    return value, action(value)
            except Exception:
                value = None
            if valid:
                if value.lower() in [v.lower() for v in valid]:
                    return value.lower()
            print(self.invalid_text('Invalid {0}'.format(input_type)))


class DiscDrive:
    '''
    Interaction with disc drive
    '''
    def __init__(self, opts):
        '''
        Setup options
        '''
        self.opts = opts

    def close(self):
        '''
        Close disc drive drawer
        '''
        if call(['eject', '-t', self.opts['device']]) != 0:
            log.error('\nFailed to close disc drive drawer')
            sys.exit(exit_codes['disc_drive'])
        if self.opts['disc_load_sleep']:
            sleep(self.opts['disc_load_sleep'])

    def open(self):
        '''
        Open disc drive drawer
        '''
        if call(['eject', self.opts['device']]) != 0:
            log.error('\nFailed to open disc drive')
            sys.exit(exit_codes['disc_drive'])


class FreeDB:
    '''
    Interaction with freedb
    '''
    def __init__(self, opts):
        '''
        Setup parameters, options
        '''
        self.opts = opts
        self.client_name, self.client_version = opts['agent'].split('/')
        disc_id = DiscID.disc_id(DiscID.open(opts['device']))
        self.info = {'disc_id': disc_id, 'record': []}

        self.input = Input(opts)
        self.header_text = TextFormat('bold', 'blue')
        self.disc_text = TextFormat('magenta')
        self.track_text = TextFormat('yellow')

        self.get_disc_info()
        self.get_track_info()
        self.set_encodings()
        self.get_preferred()

    def _present_results(self, number=None, entry=None):
        '''
        Present retrieved disc info to the user
        '''
        def present_result(number, entry):
            '''
            Print info from a single entry
            '''
            print(self.header_text('\n===== Result {0:>02} =====\n'.format(number + 1)))
            for key in entry['disc_info']:
                print(self.disc_text('{0}: {1}'.format(key, entry['disc_info'][key])))
            for track in range(self.info['disc_id'][1]):
                track_title = entry['track_info']['TTITLE{0}'.format(track)]
                print(self.track_text('{0:>02d} - {1}'.format(track + 1, track_title)))
            print(self.header_text('\n=====================\n'))

        if number is not None and entry is not None:
            present_result(number, entry)
        else:
            for num, ent in enumerate(self.info['record']):
                present_result(num, ent)

    def get_disc_info(self):
        '''
        Retrieve disc info
        '''
        status, results = 0, None
        try:
            status, results = CDDB.query(self.info['disc_id'],
                                         server_url=self.opts['freedb_mirror'],
                                         client_name=self.client_name,
                                         client_version=self.client_version)
        except IOError:
            log.error('disc_id "{0}" not found on freedb'.format(self.info['disc_id']))
        if status >= 400:
            log.error('Failed to query freedb mirror "{0}" for disc info on disc_id "{1}": '
                          'HTTP code {2}'.format(self.opts['freedb_mirror'], self.info['disc_id'][0], status))
            sys.exit(exit_codes['freedb_query'])
        if not results:
            results = [{'category': None, 'disc_id': self.info['disc_id']}]

        # A single result will not be enclosed in a list by the CDDB lib
        self.results = [results] if isinstance(results, dict) else results

    def get_track_info(self):
        '''
        Retrieve track info for each disc info result
        '''
        for result in self.results:
            if not result['category']:
                continue
            sleep(1)  # Throttle query rate
            status, track_info = CDDB.read(result['category'],
                                           result['disc_id'],
                                           server_url=self.opts['freedb_mirror'],
                                           client_name=self.client_name,
                                           client_version=self.client_version)

            if not str(status).startswith('2'):
                log.error('Failed to query freedb mirror {0} for track info: '
                              'HTTP code {1}'.format(self.opts['freedb_mirror'], status))
                sys.exit(exit_codes['freedb_query'])

            self.info['record'].append({'disc_info': result, 'track_info': track_info})

    def set_encodings(self):
        '''
        Set input character encoding for each info set
        '''
        def decode_info(coding):
            '''
            Decode disc and track info with ``coding``
            '''
            disc_info = encode_dict(entry['disc_info'], coding)
            track_info = encode_dict(entry['track_info'], coding)
            return {'disc_info': disc_info, 'track_info': track_info}

        if not len(self.info['record']):
            return  # Only the TOC is known; no disc and track titles available

        yes = ['', 'y', 'yes'] ; no = ['n', 'no']
        # This is a second order prompt loop
        for number, entry in enumerate(self.info['record']):
            suggested_encodings = suggest_encodings(entry)
            coding_prompt = 'Try an encoding (suggested: {0}): '.format(suggested_encodings)
            accepted = False
            while not accepted:
                self._present_results(number, entry)
                coding, decoded_entry = self.input(coding_prompt, 'encoding', action=decode_info)

                self._present_results(number, decoded_entry)
                accept_prompt = 'Accept source encoding ({0})? [Y/n] '.format(coding)
                accept = self.input(accept_prompt, 'answer', valid=yes + no)
                if accept in yes:
                    accepted = True
                    self.info['record'][number] = decoded_entry

    def get_preferred(self):
        '''
        Get preferred info set
        '''
        self._present_results()

        if len(self.info['record']) == 0:
            log.warning('\nNo freedb results found for disc')
            return
        elif len(self.info['record']) == 1:
            selected_entry = self.info['record'][0]
            self.info['record'][0]['preferred'] = True
        else:
            select_range = '[1-{0}]'.format(len(self.info['record']))
            prompt = 'Select a preferred result {0}: '.format(select_range)
            valid = [str(i) for i in range(1, len(self.info['record']) + 1)]
            selected_number = int(self.input(prompt, 'result', valid=valid))
            selected_entry = self.info['record'][selected_number - 1]
            selected_entry['preferred'] = True

        log.info('Selected result: {0}'.format(selected_entry['disc_info']['title']))


class Discogs:
    '''
    Interaction with discogs
    '''
    def __init__(self, opts):
        '''
        Setup options
        '''
        if not HAS_DISCOGS:
            log.error('Discogs client library is unavailable')
        else:
            self.opts = opts
            self.d = discogs_client.Client(self.opts['agent'],
                                           user_token=self.opts['token'])

    def query(self, search):
        '''
        Query discogs database
        '''
        if not HAS_DISCOGS:
            log.error('Discogs client library is unavailable')
        else:
            results = self.d.search('Mahler 9 Symphonie', type='release')
            artist = results[0].artists[0]
            releases = d.search('Herbert von Karajan', type='artist')[0].releases[1].versions[0].labels[0].releases
            for release in releases:
                print(release)


class Edit:
    '''
    Edit disc directory name and track filenames
    '''
    def __init__(self, opts, disc_info):
        '''
        Setup opts
        '''
        self.opts = opts
        self.disc_info = disc_info

        self.edit()

    def edit(self):
        '''
        Edit the disc/track info
        '''
        for number, entry in enumerate(self.disc_info['record']):
            if 'preferred' in entry:
                preferred = number
                disc_title = entry['disc_info']['title'].replace('/', '::')
                break
        # If no CDDA database records were found, there is no preferred record
        else:
            disc_title = 'Insert disc title here'
            preferred = None

        edit_buffer = ('# Do not delete or reorder lines.\n'
                       '# Lines beginning with # will be discarded.\n'
                       '# Disc title:\n'
                       + disc_title + '\n'
                       '# Track titles:\n'
                       '# Do not edit the "00 - " prefixes\n')

        if preferred is None or not range(self.disc_info['disc_id'][1]):
            edit_buffer += '01 - Add track titles here\n02 - Make sure track titles align with what is actually on the disc\n'
        else:
            for track_index in range(self.disc_info['disc_id'][1]):
                track_name = self.disc_info['record'][preferred]['track_info']['TTITLE{0}'.format(track_index)]
                new_name = '{0:02} - {1}'.format(track_index + 1, track_name.replace('/', '::'))
                edit_buffer += new_name + '\n'

        # Edit
        list_fd, list_path = tempfile.mkstemp('jmoney')
        with open(list_path, 'w') as list_file:
            list_file.write(edit_buffer)
        call([self.opts['editor'], list_path])

        # Append discID to disc name and file suffix to track names
        with open(list_path) as list_file:
            list_data = [title.strip() for title in list_file.readlines() if not title.strip().startswith('#')]
        self.disc_id = hex(self.disc_info['disc_id'][0])
        self.info = {
            'disc_title': '{}.{}'.format(list_data[0], self.disc_id),
            'track_titles': ['{}.wav'.format(title) for title in list_data[1:]],
        }
        if not self.info['disc_title']:
            log.warning('The disc title is empty')
        if len(self.info['track_titles']) != self.disc_info['disc_id'][1]:
            log.error('The number of edited titles {} is different from the number of titles indicated by CDDA database {}'.format(self.info['track_titles'], self.disc_info['disc_id'][1]))
            sys.exit(exit_codes['edit_disc'])


class CDParanoia:
    '''
    Interaction with cdparanoia
    '''
    def __init__(self, opts, disc_info, disc_id, edited_info):
        '''
        Setup options and disc information
        '''
        self.opts = opts
        self.disc_info = disc_info
        self.disc_id = disc_id
        self.edited_info = edited_info

        self.setup_dir()
        self.rip_tracks()
        self.rename_tracks()
        self.save_disc_info()

    def setup_dir(self):
        '''
        Setup disc directory
        '''
        def find_disc_id_in_library(disc_id):
            '''
            Try to find another path in the library having the same disc_id
            '''
            library = os.listdir(self.opts['audio_dir'])
            for disc_path in library:
                if disc_id in disc_path:
                    return disc_path

        other_dir = find_disc_id_in_library(self.disc_id)
        if other_dir:
            full_other_dir = os.path.join(self.opts['audio_dir'], other_dir)
            if self.opts['force']:
                log.warning('Force option enabled; duplicate library entry "{0}" ignored'.format(self.disc_id, full_other_dir))
            else:
                log.error('Disc (disc ID {0}) already in library at "{1}"'.format(self.disc_id, full_other_dir))
                sys.exit(exit_codes['disc_path'])

        self.full_disc_path = os.path.join(self.opts['audio_dir'], self.edited_info['disc_title'])
        if os.path.exists(self.full_disc_path):
            if self.opts['force']:
                log.warning('Force option enabled; overwriting destination directory, "{0}"'.format(self.full_disc_path))
                for track in os.listdir(self.full_disc_path):
                    os.remove(os.path.join(self.full_disc_path, track))
                os.rmdir(self.full_disc_path)
            else:
                log.error('Destination directory, "{0}", exists'.format(self.full_disc_path))
                sys.exit(exit_codes['disc_path'])

        if not os.path.exists(self.full_disc_path):
            os.makedirs(self.full_disc_path)
        log.info('Writing disc to {0}'.format(self.full_disc_path))

    def rip_tracks(self):
        '''
        Rip all tracks from the disc
        '''
        cmd = ['cdparanoia', '--log-debug=/dev/null',
                             '--log-summary=/dev/null',
                             '--output-wav',
                             '--batch',
                             '--force-read-speed={0}'.format(self.opts['read_speed']),
                             '--never-skip={0}'.format(self.opts['never_skip'])]
        if call(cmd, cwd=self.full_disc_path) != 0:
            log.error('\nFailed to rip disc')
            sys.exit(exit_codes['rip_disc'])

    def rename_tracks(self):
        '''
        Rename ripped tracks
        '''
        # Not all discs have a TOC track
        TOC_source = os.path.join(self.full_disc_path, 'track00.cdda.wav')
        if os.path.exists(TOC_source):
            TOC_dest = os.path.join(self.full_disc_path, '00 - disc TOC.wav')
            os.rename(TOC_source, TOC_dest)

        # Ripped tracks may be different than tracks from CDDA database
        wav_names = [wav_name for wav_name in glob.glob(os.path.join(self.full_disc_path, '?? - *.wav')) if not wav_name.startswith('00 - ')]
        if len(wav_names) == len(self.edited_info['track_titles']):
            log.error('Number of ripped tracks "{}" does not match number of track titles "{}"'.format(len(wav_names), len(self.edited_info['track_titles'])))
            sys.exit(exit_codes['rename_tracks'])

        for track_index, track_name in enumerate(self.edited_info['track_titles']):
            old_name = 'track{0:02}.cdda.wav'.format(track_index + 1)
            new_name = self.edited_info['track_titles'][track_index]
            os.rename(os.path.join(self.full_disc_path, old_name), os.path.join(self.full_disc_path, new_name))

    def save_disc_info(self):
        '''
        Save disc TOC and database record in disc directory
        '''
        if not len(self.disc_info['record']):
            log.warning('\nNo CDDA database records were found for {0}; '
                     ''.format(self.disc_info['disc_id'][0], self.full_disc_path))

        with open(os.path.join(self.full_disc_path, '00 - disc info.yaml'), 'wb') as disc_record:
            disc_record.write(yaml.dump(self.disc_info))


class FLAC:
    '''
    Flac interaction
    '''
    def __init__(self, opts, full_disc_path):
        '''
        Setup opts
        '''
        self.opts = opts
        self.full_disc_path = full_disc_path
        self.encode()

    def encode(self):
        '''
        Encode all `*.wav` files in the disc directory
        '''
        def show(disc_path):
            '''
            Report the size of disc_path
            '''
            call(['du', '-h', disc_path])

        def encode_disc(cmd, disc_path):
            '''
            Encode all the tracks from a disc
            '''
            for track_name in sorted(os.listdir(disc_path)):
                if track_name.endswith('.wav'):
                    call(cmd + [track_name], cwd=disc_path)

        cmd = ['flac', '--silent']
        if self.opts['verify_encoding']:
            cmd.append('--verify')
        if self.opts['delete_wav']:
            cmd.append('--delete-input-file')

        log.info('Compressing tracks in disc dir {}'.format(self.full_disc_path))
        encode_disc(cmd, self.full_disc_path)
        show(self.full_disc_path)


class MakeMKV:
    '''
    MakeMKV interaction
    '''
    def __init__(self, opts):
        '''
        Setup opts
        '''
        self.opts = opts

    def rip(self):
        '''
        Execute MakeMKV command line
        '''
        # TODO: Find a better construct
        if self.opts['device'] == '/dev/sr0':
            device = 'disc:0'
        elif self.opts['device'] == '/dev/sr1':
            device = 'disc:1'
        call(['makemkvcon', '--minlength=0', '--robot', '--decrypt', 'mkv', device, 'all', self.opts['video_dir']])


def get_opts():
    '''
    Setup program options
    '''
    defaults = {
        'type': 'audio',
        'quiet': False,
        'color': True,
        'verify_encoding': True,
        'force': False,
        'delete_wav': True,
        'disc_load_sleep': 29,
        'read_speed': 8,
        'never_skip': 32,
        'disc_device': '/dev/sr0',
        'freedb_mirror': 'http://freedb.freedb.org/~cddb/cddb.cgi',
        'editor': 'vim',
        'audio_dir': os.path.join(os.environ['HOME'], 'Music'),
        'script_dir': os.path.split(__file__)[0],
        'agent_file': os.path.join(os.path.split(__file__)[0], 'agent'),
        'token_file': os.path.join(os.path.split(__file__)[0], 'token'),
    }

    def block(b):
        '''
        Check whether the supplied path exists and exit if it is not a block
        device
        '''
        if os.path.exists(b):
            if stat.S_ISBLK(os.stat(b).st_mode):
                return b
            # HACK: Disc has been preripped (by `dvdbackup`, for example) to a directory
            elif os.path.isdir(b):
                return b
            else:
                raise ArgumentTypeError('{} is not a block device')
        else:
            raise ArgumentTypeError('{} does not exist')

    def new_directory(d):
        '''
        Create a new directory and error if it already exists
        '''
        return directory(d, new_directory=True)

    def directory(d, new_directory=False):
        '''
        Check whether the supplied path exists, create it if it doesn't,
        and exit if it is not a directory
        '''
        d = os.path.expanduser(os.path.expandvars(d))
        if not os.path.exists(d):
            if not os.path.islink(d):
                os.makedirs(d)
                return d
            raise ArgumentTypeError('"{}" already exists as a link'.format(d))
        elif os.path.isdir(d):
            if new_directory:
                raise ArgumentTypeError('"{}" already exists'.format(d))
            return d
        else:
            raise ArgumentTypeError('"{}" already exists as a file'.format(d))

    def parse_args():
        '''
        process invocation arguments
        '''
        desc = ('Rip an optical disc to compressed tracks or an image')
        parser = ArgumentParser(description=desc, formatter_class=ArgumentDefaultsHelpFormatter)
        parser.add_argument('--type',
                            type=str,
                            default=defaults['type'],
                            help='Type of action to execute')
        parser.add_argument('-q', '--quiet',
                            type=bool,
                            default=defaults['quiet'],
                            help='Quiet output')
        parser.add_argument('-c', '--color',
                            type=bool,
                            default=defaults['color'],
                            help='Toggle colorized output')
        parser.add_argument('-d', '--device',
                            type=block,
                            default=defaults['disc_device'],
                            help='Disc drive device name')
        parser.add_argument('-a', '--agent-file',
                            type=FileType('rb'),
                            default=defaults['agent_file'],
                            help='File containing HTTP client agent name and version conforming to RFC 1945 §3.7')
        parser.add_argument('-t', '--token-file',
                            type=FileType('rb'),
                            default=defaults['token_file'],
                            help='File containing discogs API token')
        parser.add_argument('-m', '--freedb-mirror',
                            type=str,
                            default=defaults['freedb_mirror'],
                            help='Freedb mirror URL')
        parser.add_argument('-e', '--editor',
                            type=str,
                            default=defaults['editor'],
                            help='Text editor used for updating disc/track names')
        parser.add_argument('-u', '--audio-dir',
                            type=directory,
                            default=defaults['audio_dir'],
                            help='Base directory of music library')
        parser.add_argument('--disc-load-sleep',
                            type=int,
                            default=defaults['disc_load_sleep'],
                            help='Disc drive sleep before accessing disc')
        parser.add_argument('-s', '--read-speed',
                            type=int,
                            default=defaults['read_speed'],
                            help='Disc drive read speed')
        parser.add_argument('-S', '--never-skip',
                            type=int,
                            default=defaults['never_skip'],
                            help='Reread skips until successful or this many times; if set to zero, reread indefinitely')
        parser.add_argument('--force',
                            type=bool,
                            default=defaults['force'],
                            help='Overwrite existing disc directory if it exists or ignore existing disc director(y|ies) with same disc_id')
        parser.add_argument('--verify-encoding',
                            type=bool,
                            default=defaults['verify_encoding'],
                            help='Do not verify that flac encoded files match their source wav files')
        parser.add_argument('--delete-wav',
                            type=bool,
                            default=defaults['delete_wav'],
                            help='Do not delete wave files after successful flac encoding')
        parser.add_argument('--video-dir',
                            type=new_directory,
                            help='Path to directory used by MakeMKV to store video tracks')

        return vars(parser.parse_args())

    opts = parse_args()
    opts['agent'] = opts['agent_file'].read().strip() ; del opts['agent_file']
    opts['token'] = opts['token_file'].read().strip() ; del opts['token_file']

    return opts


def main():
    '''
    Rip, encode, and name disc dir, track files
    '''
    opts = get_opts()
    log.setLevel('WARNING' if opts['quiet'] else 'INFO')
    disc_drive = DiscDrive(opts)
    disc_drive.close()

    if opts['type'] == 'audio':
        freedb = FreeDB(opts)
        edit = Edit(opts, freedb.info)
        cdparanoia = CDParanoia(opts, freedb.info, edit.disc_id, edit.info)
        flac = FLAC(opts, cdparanoia.full_disc_path)
    elif opts['type'] == 'video':
        make_mkv = MakeMKV(opts)
        make_mkv.rip()

    disc_drive.open()


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass