#!/usr/pkg/bin/python3.8

# Audio Tools, a module and set of tools for manipulating audio data
# Copyright (C) 2007-2015  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 os
import os.path
import sys
import tempfile
import subprocess
import audiotools
import audiotools.toc
import audiotools.text as _
from audiotools.decoders import SameSample

MAX_CPUS = audiotools.MAX_JOBS


def convert_to_wave(progress, audiofile, wave_filename):
    try:
        if (((audiofile.sample_rate() == 44100) and
             (audiofile.channels() == 2) and
             (audiofile.bits_per_sample() == 16))):  # already CD quality
            audiofile.convert(target_path=wave_filename,
                              target_class=audiotools.WaveAudio,
                              progress=progress)

        else:                                        # convert to CD quality
            pcm = audiotools.PCMReaderProgress(
                pcmreader=audiotools.PCMConverter(
                    audiofile.to_pcm(),
                    sample_rate=44100,
                    channels=2,
                    channel_mask=audiotools.ChannelMask.from_channels(2),
                    bits_per_sample=16),
                total_frames=(audiofile.total_frames() *
                              44100 // audiofile.sample_rate()),
                progress=progress)
            audiotools.WaveAudio.from_pcm(wave_filename, pcm)
            pcm.close()
    except audiotools.EncodingError as err:
        pass
    except KeyboardInterrupt:
        pass


def CD_reader_size(audiofile):
    """returns a (PCMReader, total_pcm_frames) of track in CD format"""

    if (((audiofile.bits_per_sample() == 16) and
         (audiofile.channels() == 2) and
         (audiofile.sample_rate() == 44100) and
         audiofile.lossless())):
        return (audiofile.to_pcm(), audiofile.total_frames())
    else:
        # if input audio isn't CD quality or lossless,
        # pull audio data into RAM so PCM frame size
        # can be calculated exactly
        pcm_frames = []
        audiotools.transfer_data(
            audiotools.PCMConverter(audiofile.to_pcm(),
                                    44100, 2, 0x3, 16).read,
            pcm_frames.append)
        return (PCMBuffer(pcm_frames, 44100, 2, 0x3, 16),
                sum(f.frames for f in pcm_frames))


class PCMBuffer(object):
    def __init__(self, pcm_frames,
                 sample_rate, channels, channel_mask, bits_per_sample):
        from collections import deque

        self.pcm_frames = deque(pcm_frames)
        self.sample_rate = sample_rate
        self.channels = channels
        self.channel_mask = channel_mask
        self.bits_per_sample = bits_per_sample

    def read(self, frames):
        if self.pcm_frames is not None:
            try:
                return self.pcm_frames.popleft()
            except IndexError:
                from audiotools.pcm import empty_framelist
                return empty_framelist(self.channels, self.bits_per_sample)
        else:
            raise ValueError()

    def close(self):
        self.pcm_frames = None


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

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

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

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

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

    parser.add_argument("-s", "--speed",
                        dest="speed",
                        default=20,
                        type=int,
                        help=_.OPT_SPEED)

    parser.add_argument("--cue",
                        dest="cuesheet",
                        metavar="FILENAME",
                        help=_.OPT_CUESHEET_TRACK2CD)

    parser.add_argument("-j", "--joint",
                        type=int,
                        default=MAX_CPUS,
                        dest="max_processes",
                        help=_.OPT_JOINT)

    parser.add_argument("filenames",
                        metavar="FILENAME",
                        nargs="+",
                        help=_.OPT_INPUT_FILENAME)

    options = parser.parse_args()

    msg = audiotools.Messenger(options.verbosity == "quiet")

    if options.max_processes < 1:
        msg.error(_.ERR_INVALID_JOINT)
        sys.exit(1)
    else:
        max_processes = options.max_processes

    write_offset = audiotools.config.getint_default("System",
                                                    "cdrom_write_offset",
                                                    0)

    audiofiles = []
    for filename in map(audiotools.Filename, options.filenames):
        try:
            audiofile = audiotools.open(str(filename))
            if audiofile.__class__.supports_to_pcm():
                audiofiles.append(audiofile)
            else:
                msg.error(_.ERR_UNSUPPORTED_TO_PCM %
                          {"filename": filename, "type": audiofile.NAME})
                sys.exit(1)
        except audiotools.UnsupportedFile:
            msg.error(_.ERR_UNSUPPORTED_FILE % (filename,))
            sys.exit(1)
        except audiotools.InvalidFile as err:
            msg.error(str(err))
            sys.exit(1)
        except IOError:
            msg.error(_.ERR_OPEN_IOERROR % (filename,))
            sys.exit(1)

    if len(audiofiles) < 1:
        msg.error(_.ERR_FILES_REQUIRED)
        sys.exit(1)

    if (((len(audiofiles) == 1) and
         (audiofiles[0].get_cuesheet() is not None))):
        # writing a single file with an embedded cuesheet
        # so extract its contents to wave/cue and call cdrdao

        if not audiotools.BIN.can_execute(audiotools.BIN['cdrdao']):
            msg.error(_.ERR_NO_CDRDAO)
            msg.info(_.ERR_GET_CDRDAO)
            sys.exit(1)

        cuesheet = audiofiles[0].get_cuesheet()

        temptoc = tempfile.NamedTemporaryFile(suffix='.toc')
        tempwav = tempfile.NamedTemporaryFile(suffix='.wav')

        audiotools.toc.write_tocfile(
            cuesheet,
            audiotools.Filename(tempwav.name).__unicode__(),
            temptoc)
        temptoc.flush()

        progress = audiotools.SingleProgressDisplay(
            msg, _.LAB_CONVERTING_FILE)

        try:
            if (((audiofiles[0].sample_rate() == 44100) and
                 (audiofiles[0].channels() == 2) and
                 (audiofiles[0].bits_per_sample() == 16))):
                # already CD quality, so no additional conversion necessary
                temptrack = audiotools.WaveAudio.from_pcm(
                    tempwav.name,
                    audiotools.PCMReaderProgress(
                        pcmreader=(audiofiles[0].to_pcm() if
                                   (write_offset == 0) else
                                   audiotools.PCMReaderWindow(
                                       audiofiles[0].to_pcm(),
                                       write_offset,
                                       audiofiles[0].total_frames())),
                        total_frames=audiofiles[0].total_frames(),
                        progress=progress.update))

            else:
                # not CD quality, so convert and adjust total frames as needed
                temptrack = audiotools.WaveAudio.from_pcm(
                    tempwav.name,
                    audiotools.PCMReaderProgress(
                        pcmreader=audiotools.PCMConverter(
                            pcmreader=(audiofiles[0].to_pcm() if
                                       (write_offset == 0) else
                                       audiotools.PCMReaderWindow(
                                           audiofiles[0].to_pcm(),
                                           write_offset,
                                           audiofiles[0].total_frames())),
                            sample_rate=44100,
                            channels=2,
                            channel_mask=0x3,
                            bits_per_sample=16),
                        total_frames=(audiofiles[0].total_frames() *
                                      44100 //
                                      audiofiles[0].sample_rate()),
                        progress=progress.update))
        except KeyboardInterrupt:
            progress.clear_rows()
            tempwav.close()
            msg.error(_.ERR_CANCELLED)
            sys.exit(1)

        progress.clear_rows()

        os.chdir(os.path.dirname(tempwav.name))
        cdrdao_args = [audiotools.BIN["cdrdao"], "write"]

        cdrdao_args.append("--device")
        cdrdao_args.append(options.dev)

        cdrdao_args.append("--speed")
        cdrdao_args.append(str(options.speed))

        cdrdao_args.append(temptoc.name)

        if options.verbosity != 'quiet':
            subprocess.call(cdrdao_args)
        else:
            devnull = open(os.devnull, 'wb')
            sub = subprocess.Popen(cdrdao_args,
                                   stdout=devnull,
                                   stderr=devnull)
            sub.wait()
            devnull.close()

        temptoc.close()
        tempwav.close()

    elif options.cuesheet is not None:
        # writing tracks with a cuesheet,
        # so combine them into a single wave/cue and call cdrdao

        if not audiotools.BIN.can_execute(audiotools.BIN['cdrdao']):
            msg.error(_.ERR_NO_CDRDAO)
            msg.info(_.ERR_GET_CDRDAO)
            sys.exit(1)

        if len({f.sample_rate() for f in audiofiles}) != 1:
            msg.error(_.ERR_SAMPLE_RATE_MISMATCH)
            sys.exit(1)

        if len({f.channels() for f in audiofiles}) != 1:
            msg.error(_.ERR_CHANNEL_COUNT_MISMATCH)
            sys.exit(1)

        if len({f.bits_per_sample() for f in audiofiles}) != 1:
            msg.error(_.ERR_BPS_MISMATCH)
            sys.exit(1)

        try:
            cuesheet = audiotools.read_sheet(options.cuesheet)
        except audiotools.SheetException as err:
            msg.error(err)
            sys.exit(1)

        # ensure sheet is formatted for CD images
        if not cuesheet.image_formatted():
            msg.error(_.ERR_CUE_INVALID_FORMAT)
            sys.exit(1)

        # ensure sheet fits the tracks it's given
        if len(cuesheet) != len(audiofiles):
            msg.error(_.ERR_CUE_INSUFFICIENT_TRACKS)
            sys.exit(1)

        for (input_track, cuesheet_track) in zip(audiofiles,
                                                 cuesheet.track_numbers()):
            track_length = cuesheet.track_length(cuesheet_track)
            if ((track_length is not None) and
                (track_length != input_track.seconds_length())):
                msg.error(_.ERR_CUE_LENGTH_MISMATCH % (cuesheet_track))
                sys.exit(1)

        tempcue = tempfile.NamedTemporaryFile(suffix='.cue')
        tempwav = tempfile.NamedTemporaryFile(suffix='.wav')

        audiotools.toc.write_tocfile(
            cuesheet,
            audiotools.Filename(tempwav.name).__unicode__(),
            tempcue)
        tempcue.flush()

        input_frames = (int(audiofiles[0].sample_rate() * cuesheet.pre_gap()) +
                        sum([af.total_frames() for af in audiofiles]))

        if (((audiofiles[0].sample_rate() == 44100) and
             (audiofiles[0].channels() == 2) and
             (audiofiles[0].bits_per_sample() == 16))):
            pcmreader = audiotools.PCMCat(
                [SameSample(
                    sample=0,
                    total_pcm_frames=int(44100 * cuesheet.pre_gap()),
                    sample_rate=44100,
                    channels=2,
                    channel_mask=0x3,
                    bits_per_sample=16)] +
                [af.to_pcm() for af in audiofiles])
            output_frames = input_frames
        else:
            # this presumes a cuesheet and non-CD audio
            # though theoretically possible, it's difficult to
            # envision a case in which this happens
            pcmreader = audiotools.PCMConverter(
                pcmreader=audiotools.PCMCat(
                    [SameSample(
                        sample=0,
                        total_pcm_frames=int(audiofiles[0].sample_rate() *
                                             cuesheet.pre_gap()),
                        channels=audiofiles[0].channels(),
                        channel_mask=int(audiofiles[0].channel_mask()),
                        bits_per_sample=audiofiles[0].bits_per_sample())] +
                    [af.to_pcm() for af in audiofiles]),
                sample_rate=44100,
                channels=2,
                channel_mask=0x3,
                bits_per_sample=16)
            output_frames = (input_frames * 44100 //
                             audiofiles[0].sample_rate())

        progress = audiotools.SingleProgressDisplay(
            msg, _.LAB_CONVERTING_FILE)

        try:
            if write_offset == 0:
                temptrack = audiotools.WaveAudio.from_pcm(
                    tempwav.name,
                    audiotools.PCMReaderProgress(
                        pcmreader=pcmreader,
                        total_frames=output_frames,
                        progress=progress.update))
            else:
                temptrack = audiotools.WaveAudio.from_pcm(
                    tempwav.name,
                    audiotools.PCMReaderProgress(
                        pcmreader=audiotools.PCMReaderWindow(pcmreader,
                                                             write_offset,
                                                             input_frames),
                        total_frames=output_frames,
                        progress=progress.update))
        except KeyboardInterrupt:
            progress.clear_rows()
            tempwav.close()
            msg.error(_.ERR_CANCELLED)
            sys.exit(1)

        progress.clear_rows()

        os.chdir(os.path.dirname(tempwav.name))

        cdrdao_args = [audiotools.BIN["cdrdao"], "write",
                       "--device", options.dev,
                       "--speed", str(options.speed),
                       tempcue.name]

        if options.verbosity != 'quiet':
            subprocess.call(cdrdao_args)
        else:
            devnull = open(os.devnull, 'wb')
            sub = subprocess.Popen(cdrdao_args,
                                   stdout=devnull,
                                   stderr=devnull)
            sub.wait()
            devnull.close()

        tempcue.close()
        tempwav.close()
    else:
        # writing tracks without a cuesheet,
        # so extract them to waves and call cdrecord

        if not audiotools.BIN.can_execute(audiotools.BIN['cdrecord']):
            msg.error(_.ERR_NO_CDRECORD)
            msg.info(_.ERR_GET_CDRECORD)
            sys.exit(1)

        exec_args = [audiotools.BIN['cdrecord'],
                     "-pad",
                     "-speed", str(options.speed),
                     "-dev", options.dev,
                     "-dao", "-audio"]

        temp_pool = []
        wave_files = []

        if write_offset == 0:
            queue = audiotools.ExecProgressQueue(msg)

            for audiofile in audiofiles:
                if isinstance(audiofile, audiotools.WaveAudio):
                    wave_files.append(audiofile.filename)
                else:
                    filename = audiotools.Filename(audiofile.filename)
                    f = tempfile.mkstemp(suffix='.wav')
                    temp_pool.append(f)
                    wave_files.append(f[1])
                    queue.execute(
                        function=convert_to_wave,
                        progress_text=filename.__unicode__(),
                        completion_output=_.LAB_TRACK2CD_CONVERTED % {
                            "filename": filename},
                        audiofile=audiofile,
                        wave_filename=f[1])

            try:
                queue.run(max_processes)
            except KeyboardInterrupt:
                msg.error(_.ERR_CANCELLED)
                for (fd, f) in temp_pool:
                    os.close(fd)
                    os.unlink(f)
                sys.exit(1)
        else:
            # a list of (PCMReader, total_pcm_frames) pairs
            try:
                lossless_tracks = [CD_reader_size(t) for t in audiofiles]

                realigned_stream = audiotools.BufferedPCMReader(
                    audiotools.PCMReaderWindow(
                        audiotools.PCMCat([t[0] for t in lossless_tracks]),
                        write_offset,
                        sum([t[1] for t in lossless_tracks])))

                for audiofile in audiofiles:
                    filename = audiotools.Filename(audiofile.filename)

                    progress = audiotools.SingleProgressDisplay(
                        msg, filename.__unicode__())

                    f = tempfile.mkstemp(suffix=".wav")
                    temp_pool.append(f)
                    wave = audiotools.WaveAudio.from_pcm(
                        f[1],
                        audiotools.PCMReaderProgress(
                            audiotools.LimitedPCMReader(
                                realigned_stream, audiofile.total_frames()),
                            audiofile.total_frames(),
                            progress.update))
                    wave_files.append(wave.filename)
                    progress.clear_rows()
                    msg.info(_.LAB_TRACK2CD_CONVERTED % {"filename": filename})
            except KeyboardInterrupt:
                progress.clear_rows()
                msg.error(_.ERR_CANCELLED)
                for (fd, f) in temp_pool:
                    os.close(fd)
                    os.unlink(f)
                sys.exit(1)

        try:
            for wave in wave_files:
                audiotools.open(wave).verify()
        except (audiotools.UnsupportedFile,
                audiotools.InvalidFile,
                IOError):
            msg.error(_.ERR_TRACK2CD_INVALIDFILE)
            sys.exit(1)

        exec_args += wave_files

        if options.verbosity != 'quiet':
            subprocess.call(exec_args)
        else:
            devnull = open(os.devnull, 'wb')
            sub = subprocess.Popen(exec_args,
                                   stdout=devnull,
                                   stderr=devnull)
            sub.wait()
            devnull.close()

        for (fd, f) in temp_pool:
            os.close(fd)
            os.unlink(f)
