#!/usr/bin/python

# Audio Tools, a module and set of tools for manipulating audio data
# Copyright (C) 2007-2014  Brian Langenberger

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA


import sys
import os
import audiotools
import audiotools.ui
import audiotools.accuraterip
from audiotools.cdio import CDDAReader
import termios
import audiotools.text as _

PREVIOUS_TRACK_FRAMES = (5880 // 2)
NEXT_TRACK_FRAMES = (5880 // 2)


class AccurateRipReader(object):
    def __init__(self, pcmreader, total_pcm_frames, is_first, is_last):
        """pcmreader is a PCMReader object to wrap around

        total_pcm_frames is the length of pcmreader,
        not including previous and next track frames

        is_first and is_last indicate the track's position in the stream"""

        self.__read__ = pcmreader.read
        self.__close__ = pcmreader.close

        self.__checksummer__ = audiotools.accuraterip.Checksum(
            total_pcm_frames=total_pcm_frames,
            sample_rate=pcmreader.sample_rate,
            is_first=is_first,
            is_last=is_last,
            pcm_frame_range=PREVIOUS_TRACK_FRAMES + 1 + NEXT_TRACK_FRAMES,
            accurateripv2_offset=PREVIOUS_TRACK_FRAMES)

        self.sample_rate = pcmreader.sample_rate
        self.channels = pcmreader.channels
        self.channel_mask = pcmreader.channel_mask
        self.bits_per_sample = pcmreader.bits_per_sample

    def read(self, pcm_frames):
        frame = self.__read__(pcm_frames)
        self.__checksummer__.update(frame)
        return frame

    def close(self):
        self.__close__()

    def checksums_v1(self):
        return self.__checksummer__.checksums_v1()

    def checksums_v2(self):
        return [self.__checksummer__.checksum_v2()]


if (__name__ == '__main__'):
    import argparse

    parser = argparse.ArgumentParser(description=_.DESCRIPTION_CD2TRACK)

    parser.add_argument("--version",
                        action="version",
                        version="Python Audio Tools %s" % (audiotools.VERSION))

    parser.add_argument("-I", "--interactive",
                        action="store_true",
                        default=False,
                        dest="interactive",
                        help=_.OPT_INTERACTIVE_OPTIONS)

    parser.add_argument("-V", "--verbose",
                        dest="verbosity",
                        choices=audiotools.VERBOSITY_LEVELS,
                        default=audiotools.DEFAULT_VERBOSITY,
                        help=_.OPT_VERBOSE)

    parser.add_argument("-c", "--cdrom",
                        dest="cdrom",
                        default=audiotools.DEFAULT_CDROM)

    parser.add_argument("-s", "--speed",
                        type=int,
                        dest="speed")

    conversion = parser.add_argument_group(_.OPT_CAT_EXTRACTION)

    conversion.add_argument("-t", "--type",
                            dest="type",
                            choices=sorted(audiotools.TYPE_MAP.keys() +
                                           ["help"]),
                            default=audiotools.DEFAULT_TYPE,
                            help=_.OPT_TYPE)

    conversion.add_argument("-q", "--quality",
                            dest="quality",
                            help=_.OPT_QUALITY)

    conversion.add_argument("-d", "--dir",
                            dest="dir",
                            default=".",
                            help=_.OPT_DIR)

    conversion.add_argument("--format",
                            default=None,
                            dest="format",
                            help=_.OPT_FORMAT)

    lookup = parser.add_argument_group(_.OPT_CAT_CD_LOOKUP)

    lookup.add_argument("-M", "--metadata-lookup",
                        action="store_true",
                        default=False,
                        dest="metadata_lookup",
                        help=_.OPT_METADATA_LOOKUP)

    lookup.add_argument("--musicbrainz-server",
                        dest="musicbrainz_server",
                        default=audiotools.MUSICBRAINZ_SERVER,
                        metavar="HOSTNAME")

    lookup.add_argument("--musicbrainz-port",
                        type=int,
                        dest="musicbrainz_port",
                        default=audiotools.MUSICBRAINZ_PORT,
                        metavar="PORT")

    lookup.add_argument("--no-musicbrainz",
                        action="store_false",
                        dest="use_musicbrainz",
                        default=audiotools.MUSICBRAINZ_SERVICE,
                        help=_.OPT_NO_MUSICBRAINZ)

    lookup.add_argument("--freedb-server",
                        dest="freedb_server",
                        default=audiotools.FREEDB_SERVER,
                        metavar="HOSTNAME")

    lookup.add_argument("--freedb-port",
                        type=int,
                        dest="freedb_port",
                        default=audiotools.FREEDB_PORT,
                        metavar="PORT")

    lookup.add_argument("--no-freedb",
                        action="store_false",
                        dest="use_freedb",
                        default=audiotools.FREEDB_SERVICE,
                        help=_.OPT_NO_FREEDB)

    lookup.add_argument("-D", "--default",
                        dest="use_default",
                        action="store_true",
                        default=False,
                        help=_.OPT_DEFAULT)

    metadata = parser.add_argument_group(_.OPT_CAT_METADATA)

    metadata.add_argument("--album-number",
                          dest="album_number",
                          type=int,
                          default=0,
                          help=_.OPT_ALBUM_NUMBER)

    metadata.add_argument("--album-total",
                          dest="album_total",
                          type=int,
                          default=0,
                          help=_.OPT_ALBUM_TOTAL)

    # if adding ReplayGain is a lossless process
    # (i.e. added as tags rather than modifying track data)
    # add_replay_gain should default to True
    # if not, add_replay_gain should default to False
    # which is which depends on the track type
    metadata.add_argument("--replay-gain",
                          action="store_true",
                          default=None,
                          dest="add_replay_gain",
                          help=_.OPT_REPLAY_GAIN)

    metadata.add_argument("--no-replay-gain",
                          action="store_false",
                          default=None,
                          dest="add_replay_gain",
                          help=_.OPT_NO_REPLAY_GAIN)

    parser.add_argument("tracks",
                        type=int,
                        metavar="TRACK",
                        nargs="*",
                        help=_.OPT_TRACK_INDEX)

    options = parser.parse_args()

    msg = audiotools.Messenger("cd2track", options)

    # ensure interactive mode is available, if selected
    if (options.interactive and (not audiotools.ui.AVAILABLE)):
        audiotools.ui.not_available_message(msg)
        sys.exit(1)

    # get the default AudioFile class we are converted to
    if (options.type == 'help'):
        audiotools.ui.show_available_formats(msg)
        sys.exit(0)
    else:
        AudioType = audiotools.TYPE_MAP[options.type]

    # ensure the selected compression is compatible with that class
    if (options.quality == 'help'):
        audiotools.ui.show_available_qualities(msg, AudioType)
        sys.exit(0)
    elif (options.quality is None):
        options.quality = audiotools.__default_quality__(AudioType.NAME)
    elif (options.quality not in AudioType.COMPRESSION_MODES):
        msg.error(_.ERR_UNSUPPORTED_COMPRESSION_MODE %
                  {"quality": options.quality,
                   "type": AudioType.NAME})
        sys.exit(1)

    quality = options.quality
    base_directory = options.dir

    try:
        cddareader = CDDAReader(options.cdrom, True)
        track_offsets = cddareader.track_offsets
        track_lengths = cddareader.track_lengths
    except (IOError, ValueError) as err:
        msg.error(unicode(err) + _.ERR_INVALID_CDDA)
        sys.exit(-1)

    # only apply read_offset if cddareader is a physical drive
    if (not cddareader.is_cd_image):
        try:
            read_offset = audiotools.config.getint_default(
                "System", "cdrom_read_offset", 0)
        except ValueError:
            read_offset = 0
    else:
        read_offset = 0

    # set speed, if any
    if (options.speed is not None):
        cddareader.set_speed(options.speed)

    # use CDDA object to query metadata services for metadata choices
    try:
        metadata_choices = audiotools.cddareader_metadata_lookup(
            cddareader=cddareader,
            musicbrainz_server=options.musicbrainz_server,
            musicbrainz_port=options.musicbrainz_port,
            freedb_server=options.freedb_server,
            freedb_port=options.freedb_port,
            use_musicbrainz=options.use_musicbrainz,
            use_freedb=options.use_freedb)
    except KeyboardInterrupt:
        msg.ansi_clearscreen()
        msg.error(_.ERR_CANCELLED)
        sys.exit(1)

    # update MetaData with command-line album-number/total, if given
    if (options.album_number != 0):
        for c in metadata_choices:
            for m in c:
                m.album_number = options.album_number

    if (options.album_total != 0):
        for c in metadata_choices:
            for m in c:
                m.album_total = options.album_total

    # use CDDAReader object to perform AccurateRip lookup
    try:
        ar_result = audiotools.accuraterip.perform_lookup(
            audiotools.accuraterip.DiscID.from_cddareader(cddareader))
    except KeyboardInterrupt:
        msg.ansi_clearline()
        msg.error(_.ERR_CANCELLED)
        sys.exit(1)

    # determine tracks to be ripped
    if (len(options.tracks) == 0):
        tracks_to_rip = list(sorted(track_offsets.keys()))
    else:
        tracks_to_rip = [t for t in options.tracks if
                         t in track_offsets.keys()]

    if (len(tracks_to_rip) == 0):
        # no tracks selected to rip, so do nothing
        sys.exit(0)

    # decide which metadata and output options use when extracting tracks
    if (options.interactive):
        # pick options using interactive widget
        output_widget = audiotools.ui.OutputFiller(
            track_labels=[_.LAB_TRACK_X_OF_Y % (i, max(track_offsets.keys()))
                          for i in tracks_to_rip],
            metadata_choices=[[c[i - 1] for i in tracks_to_rip]
                              for c in metadata_choices],
            input_filenames=[audiotools.Filename("track%2.2d.cdda.wav" % (i))
                             for i in tracks_to_rip],
            output_directory=options.dir,
            format_string=(options.format if
                           (options.format is not None) else
                           audiotools.FILENAME_FORMAT),
            output_class=AudioType,
            quality=quality,
            completion_label=_.LAB_CD2TRACK_APPLY)

        loop = audiotools.ui.urwid.MainLoop(
            output_widget,
            audiotools.ui.style(),
            unhandled_input=output_widget.handle_text,
            pop_ups=True)
        try:
            loop.run()
            msg.ansi_clearscreen()
        except (termios.error, IOError):
            msg.error(_.ERR_TERMIOS_ERROR)
            msg.info(_.ERR_TERMIOS_SUGGESTION)
            msg.info(audiotools.ui.xargs_suggestion(sys.argv))
            sys.exit(1)

        if (not output_widget.cancelled()):
            output_tracks = list(output_widget.output_tracks())
        else:
            sys.exit(0)
    else:
        # pick options without using GUI
        try:
            output_tracks = list(
                audiotools.ui.process_output_options(
                    metadata_choices=[[c[i - 1] for i in tracks_to_rip]
                                      for c in metadata_choices],
                    input_filenames=[
                        audiotools.Filename("track%2.2d.cdda.wav" % (i))
                        for i in tracks_to_rip],
                    output_directory=options.dir,
                    format_string=options.format,
                    output_class=AudioType,
                    quality=options.quality,
                    msg=msg,
                    use_default=options.use_default))
        except audiotools.UnsupportedTracknameField as err:
            err.error_msg(msg)
            sys.exit(1)
        except (audiotools.InvalidFilenameFormat,
                audiotools.OutputFileIsInput,
                audiotools.DuplicateOutputFile) as err:
            msg.error(unicode(err))
            sys.exit(1)

    # perform actual ripping of tracks from CDDA
    encoded = []
    rip_log = {}
    accuraterip_log_v1 = {}
    accuraterip_log_v2 = {}

    for (track_number,
         index,
         (output_class,
          output_filename,
          output_quality,
          output_metadata)) in zip(tracks_to_rip,
                                   range(1, len(tracks_to_rip) + 1),
                                   output_tracks):
        cddareader.reset_log()
        track_offset = (track_offsets[track_number] +
                        read_offset -
                        PREVIOUS_TRACK_FRAMES)
        track_length = track_lengths[track_number]

        # seek to indicated starting offset
        if (track_offset > 0):
            seeked_offset = cddareader.seek(track_offset)
        else:
            seeked_offset = cddareader.seek(0)

        # make leading directories, if necessary
        try:
            audiotools.make_dirs(str(output_filename))
        except OSError as err:
            msg.os_error(err)
            sys.exit(1)

        progress = audiotools.SingleProgressDisplay(
            msg, unicode(output_filename))

        track_data = audiotools.PCMReaderWindow(
            cddareader,
            track_offset - seeked_offset,
            PREVIOUS_TRACK_FRAMES + track_length + NEXT_TRACK_FRAMES)
        # ensure closing track doesn't close whole CD
        track_data.close = lambda self: None

        accuraterip = AccurateRipReader(
            track_data,
            track_length,
            track_number == min(track_offsets.keys()),
            track_number == max(track_offsets.keys()))

        try:
            track = output_class.from_pcm(
                str(output_filename),
                audiotools.PCMReaderProgress(
                    audiotools.PCMReaderWindow(
                        accuraterip,
                        PREVIOUS_TRACK_FRAMES,
                        track_length),
                    track_length,
                    progress.update),
                output_quality,
                total_pcm_frames=track_length)
            encoded.append(track)

            # since the inner PCMReaderWindow only outputs part
            # of the accuraterip reader, we need to ensure
            # anything left over in accuraterip gets processed also
            audiotools.transfer_framelist_data(accuraterip, lambda s: None)
        except audiotools.EncodingError as err:
            progress.clear_rows()
            msg.error(_.ERR_ENCODING_ERROR % (output_filename,))
            sys.exit(1)
        except KeyboardInterrupt:
            progress.clear_rows()
            try:
                os.unlink(str(output_filename))
            except OSError:
                pass
            msg.error(_.ERR_CANCELLED)
            sys.exit(1)

        track.set_metadata(output_metadata)
        progress.clear_rows()

        rip_log[track_number] = cddareader.log()
        accuraterip_log_v1[track_number] = accuraterip.checksums_v1()
        accuraterip_log_v2[track_number] = accuraterip.checksums_v2()

        msg.info(
            audiotools.output_progress(
                _.LAB_CD2TRACK_PROGRESS %
                {"track_number": track_number,
                 "filename": output_filename},
                index, len(tracks_to_rip)))

    # add ReplayGain to ripped tracks, if necessary
    if (output_class.supports_replay_gain() and
        (options.add_replay_gain if options.add_replay_gain is not None else
         audiotools.ADD_REPLAYGAIN)):
        rg_progress = audiotools.ReplayGainProgressDisplay(msg)
        rg_progress.initial_message()
        try:
            # all audio files must belong to the same album, by definition
            audiotools.add_replay_gain(encoded, rg_progress.update)
        except ValueError as err:
            rg_progress.clear_rows()
            msg.error(unicode(err))
            sys.exit(1)
        except KeyboardInterrupt:
            rg_progress.clear_rows()
            msg.error(_.ERR_CANCELLED)
            sys.exit(1)
        rg_progress.final_message()

    # display ripping log
    msg.output(_.LAB_CD2TRACK_LOG)

    table = audiotools.output_table()

    row = table.row()
    extract_headers = [u"track",
                       u"error",
                       u"skip",
                       u"atom",
                       u"edge",
                       u"drop",
                       u"dup",
                       u"drift",
                       _.LAB_TRACKVERIFY_AR_CHECKSUM,
                       _.LAB_TRACKVERIFY_AR_CHECKSUM_VER,
                       _.LAB_TRACKVERIFY_AR_OFFSET,
                       _.LAB_TRACKVERIFY_AR_CONFIDENCE]
    extract_keys = ["readrr",
                    "skip",
                    "fixup_atom",
                    "fixup_edge",
                    "fixup_dropped",
                    "fixup_duped",
                    "drift"]

    dividers = []

    row.add_column(extract_headers[0])
    dividers.append(_.DIV)
    for header in extract_headers[1:]:
        row.add_column(u" ")
        row.add_column(header, "right")
        dividers.append(u" ")
        dividers.append(_.DIV)

    table.divider_row(dividers)

    for track_number in tracks_to_rip:
        row = table.row()
        row.add_column(unicode(track_number), "right")
        log = rip_log[track_number]
        for key in extract_keys:
            row.add_column(u" ")
            row.add_column(unicode(log.get(key, 0)), "right")

        (checksum_v2,
         confidence_v2,
         offset_v2) = audiotools.accuraterip.match_offset(
            ar_result.get(track_number, []),
            accuraterip_log_v2[track_number],
            0)

        (checksum_v1,
         confidence_v1,
         offset_v1) = audiotools.accuraterip.match_offset(
            ar_result.get(track_number, []),
            accuraterip_log_v1[track_number],
            -PREVIOUS_TRACK_FRAMES)

        if (len(ar_result.get(track_number, [])) == 0):
            checksum = checksum_v2
            version = 2
            offset = offset_v2
            confidence = _.LAB_TRACKVERIFY_ACCURATERIP_NOTFOUND
        elif (confidence_v2 is not None):
            checksum = checksum_v2
            version = 2
            offset = offset_v2
            confidence = unicode(confidence_v2)
        elif (confidence_v1 is not None):
            checksum = checksum_v1
            version = 1
            offset = offset_v1
            confidence = unicode(confidence_v1)
        else:
            checksum = checksum_v2
            version = 2
            offset = offset_v2
            confidence = _.LAB_TRACKVERIFY_ACCURATERIP_MISMATCH

        # AccurateRip columns
        row.add_column(u"")
        row.add_column(u"%8.8X" % (checksum), "right")
        row.add_column(u"")
        row.add_column(u"v%d" % (version), "right")
        row.add_column(u"")
        row.add_column(unicode(offset), "right")
        row.add_column(u"")
        row.add_column(confidence, "right")

    for row in table.format(msg.output_isatty()):
        msg.output(row)
