#! /usr/bin/env python2
#
#    ReplayGain Album Analysis using gstreamer rganalysis element
#    Copyright (C) 2005,2007,2009  Michael Urman
#                         2012,14  Nick Boultbee
#                            2013  Christoph Reiter
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of version 2 of the GNU General Public License as
#    published by the Free Software Foundation.
#

from gi.repository import Gtk
from gi.repository import GObject
from gi.repository import Pango
from gi.repository import Gst
from gi.repository import GLib
from quodlibet import print_d
from quodlibet.plugins import PluginConfigMixin

from quodlibet.browsers.collection.models import EMPTY

from quodlibet.qltk.views import HintedTreeView
from quodlibet.plugins.songsmenu import SongsMenuPlugin
from quodlibet.util import cached_property

__all__ = ['ReplayGain']


def get_num_threads():
    # multiprocessing is >= 2.6.
    # Default to 2 threads if cpu_count isn't implemented for the current arch
    # or multiprocessing isn't available
    try:
        import multiprocessing
        threads = multiprocessing.cpu_count()
    except (ImportError, NotImplementedError):
        threads = 2
    return threads


class UpdateMode(object):
    """Enum-like class for update strategies"""
    ALWAYS = "always"
    ALBUM_MISSING = "album_tags_missing"
    ANY_MISSING = "any_tags_missing"


class RGAlbum(object):
    def __init__(self, rg_songs, process_mode):
        self.songs = rg_songs
        self.gain = None
        self.peak = None
        self.__should_process = None
        self.__process_mode = process_mode

    @property
    def progress(self):
        all_ = 0.0
        done = 0.0
        for song in self.songs:
            all_ += song.length
            done += song.length * song.progress

        try:
            return max(min(done / all_, 1.0), 0.0)
        except ZeroDivisionError:
            return 0.0

    @property
    def done(self):
        for song in self.songs:
            if not song.done:
                return False
        return True

    @property
    def title(self):
        if not self.songs:
            return ""
        # It's ok - any() + generator is short-cut-logic-friendly
        if not any(rgs.song("album") for rgs in self.songs):
            return "(%s)" % EMPTY
        return self.songs[0].song.comma('~artist~album')

    @property
    def error(self):
        for song in self.songs:
            if song.error:
                return True
        return False

    def write(self):
        # Don't write incomplete data
        if not self.done:
            return

        for song in self.songs:
            song._write(self.gain, self.peak)

    @classmethod
    def from_songs(cls, songs, process_mode=UpdateMode.ALWAYS):
        return RGAlbum([RGSong(s) for s in songs], process_mode)

    @cached_property
    def should_process(self):
        """Returns true if the album needs analysis, according to prefs"""
        mode = self.__process_mode
        if mode == UpdateMode.ALWAYS:
            return True
        elif mode == UpdateMode.ANY_MISSING:
            return not all([s.has_all_rg_tags for s in self.songs])
        elif mode == UpdateMode.ALBUM_MISSING:
            return not all([s.album_gain for s in self.songs])
        else:
            print_w("Invalid setting for update mode: " + mode)
            # Safest to re-process probably.
            return True


class RGSong(object):
    def __init__(self, song):
        self.song = song
        self.error = False
        self.gain = None
        self.peak = None
        self.progress = 0.0
        self.done = False
        # TODO: support prefs for not overwriting individual existing tags
        #       e.g. to re-run over entire library but keeping files untouched
        self.overwrite_existing = True

    def _write(self, album_gain, album_peak):
        if self.error or not self.done:
            return
        song = self.song

        def write_to_song(tag, pattern, value):
            if value is None or value == "":
                return
            existing = song(tag, None)
            if existing and not self.overwrite_existing:
                print_d("Not overwriting existing tag %s (=%s) for %s"
                        % (tag, existing, self.song("~filename")))
                return
            song[tag] = pattern % value

        write_to_song('replaygain_track_gain', '%.2f dB', self.gain)
        write_to_song('replaygain_track_peak', '%.4f', self.peak)
        write_to_song('replaygain_album_gain', '%.2f dB', album_gain)
        write_to_song('replaygain_album_peak', '%.4f', album_peak)

    @property
    def title(self):
        return self.song('~tracknumber~title~version')

    @property
    def filename(self):
        return self.song("~filename")

    @property
    def length(self):
        return self.song("~#length")

    def _get_rg_tag(self, suffix):
        ret = self.song("~#replaygain_%s" % suffix)
        return None if ret == "" else ret

    @property
    def track_gain(self):
        return self._get_rg_tag("track_gain")

    @property
    def album_gain(self):
        return self._get_rg_tag("album_gain")

    @property
    def track_peak(self):
        return self._get_rg_tag('track_peak')

    @property
    def album_peak(self):
        return self._get_rg_tag('album_peak')

    @property
    def has_track_tags(self):
        return not (self.track_gain is None or self.track_peak is None)

    @property
    def has_album_tags(self):
        return not (self.album_gain is None or self.album_peak is None)

    @property
    def has_all_rg_tags(self):
        return self.has_track_tags and self.has_album_tags

    def __str__(self):
        vals = {k: self._get_rg_tag(k)
                for k in 'track_gain album_gain album_peak track_peak'.split()}
        return "<Song=%s RG data=%s>" % (self.song, vals)


class ReplayGainPipeline(GObject.Object):

    __gsignals__ = {
        # done(self, album)
        'done': (GObject.SignalFlags.RUN_LAST, None, (object,)),
        # update(self, album, song)
        'update': (GObject.SignalFlags.RUN_LAST, None,
                   (object, object,)),
    }

    def __init__(self):
        super(ReplayGainPipeline, self).__init__()

        self._current = None
        self._setup_pipe()

    def _setup_pipe(self):
        # gst pipeline for replay gain analysis:
        # filesrc!decodebin!audioconvert!audioresample!rganalysis!fakesink
        self.pipe = Gst.Pipeline()
        self.filesrc = Gst.ElementFactory.make("filesrc", "source")
        self.pipe.add(self.filesrc)

        self.decode = Gst.ElementFactory.make("decodebin", "decode")

        def new_decoded_pad(dbin, pad):
            pad.link(self.convert.get_static_pad("sink"))

        def sort_decoders(decode, pad, caps, factories):
            def set_prio(x):
                i, f = x
                i = {"mad": -1, "mpg123audiodec": -2}.get(f.get_name(), i)
                return (i, f)
            return zip(*sorted(map(set_prio, enumerate(factories))))[1]

        self.decode.connect("autoplug-sort", sort_decoders)

        self.decode.connect("pad-added", new_decoded_pad)
        self.pipe.add(self.decode)
        self.filesrc.link(self.decode)

        self.convert = Gst.ElementFactory.make("audioconvert", "convert")
        self.pipe.add(self.convert)

        self.resample = Gst.ElementFactory.make("audioresample", "resample")
        self.pipe.add(self.resample)
        self.convert.link(self.resample)

        self.analysis = Gst.ElementFactory.make("rganalysis", "analysis")
        self.pipe.add(self.analysis)
        self.resample.link(self.analysis)

        self.sink = Gst.ElementFactory.make("fakesink", "sink")
        self.pipe.add(self.sink)
        self.analysis.link(self.sink)

        self.bus = bus = self.pipe.get_bus()
        bus.add_signal_watch()
        bus.connect("message", self._bus_message)

    def request_update(self):
        if not self._current:
            return

        ok, p = self.pipe.query_position(Gst.Format.TIME)
        if ok:
            length = self._current.length
            try:
                progress = float(p / Gst.SECOND) / length
            except ZeroDivisionError:
                progress = 0.0
            progress = max(min(progress, 1.0), 0.0)
            self._current.progress = progress
            self._emit_update()

    def _emit_update(self):
        self.emit("update", self._album, self._current)

    def start(self, album):
        self._album = album
        self._songs = list(album.songs)
        self._done = []
        self._next_song(first=True)

    def quit(self):
        self.bus.remove_signal_watch()
        self.pipe.set_state(Gst.State.NULL)

    def _next_song(self, first=False):
        if self._current:
            self._current.progress = 1.0
            self._current.done = True
            self._emit_update()
            self._done.append(self._current)
            self._current = None

        if not self._songs:
            self.pipe.set_state(Gst.State.NULL)
            self.emit("done", self._album)
            return

        if first:
            self.analysis.set_property("num-tracks", len(self._songs))
        else:
            self.analysis.set_locked_state(True)
            self.pipe.set_state(Gst.State.NULL)

        self._current = self._songs.pop(0)
        self.filesrc.set_property("location", self._current.filename)
        if not first:
            # flush, so the element takes new data after EOS
            pad = self.analysis.get_static_pad("src")
            pad.send_event(Gst.Event.new_flush_start())
            pad.send_event(Gst.Event.new_flush_stop(True))
            self.analysis.set_locked_state(False)
        self.pipe.set_state(Gst.State.PLAYING)

    def _bus_message(self, bus, message):
        if message.type == Gst.MessageType.TAG:
            tags = message.parse_tag()
            ok, value = tags.get_double(Gst.TAG_TRACK_GAIN)
            if ok:
                self._current.gain = value
            ok, value = tags.get_double(Gst.TAG_TRACK_PEAK)
            if ok:
                self._current.peak = value
            ok, value = tags.get_double(Gst.TAG_ALBUM_GAIN)
            if ok:
                self._album.gain = value
            ok, value = tags.get_double(Gst.TAG_ALBUM_PEAK)
            if ok:
                self._album.peak = value
            self._emit_update()
        elif message.type == Gst.MessageType.EOS:
            self._next_song()
        elif message.type == Gst.MessageType.ERROR:
            gerror, debug = message.parse_error()
            if gerror:
                print_e(gerror.message)
            print_e(debug)
            self._current.error = True
            self._next_song()


class RGDialog(Gtk.Dialog):

    def __init__(self, albums, parent, process_mode):
        super(RGDialog, self).__init__(
            title=_('ReplayGain Analyzer'), parent=parent,
            buttons=(Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL,
                     Gtk.STOCK_SAVE, Gtk.ResponseType.OK)
        )

        self.process_mode = process_mode
        self.set_default_size(600, 400)
        self.set_border_width(6)

        hbox = Gtk.HBox(spacing=6)
        info = Gtk.Label()
        hbox.pack_start(info, True, True, 0)
        self.vbox.pack_start(hbox, False, False, 6)

        swin = Gtk.ScrolledWindow()
        swin.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        swin.set_shadow_type(Gtk.ShadowType.IN)

        self.vbox.pack_start(swin, True, True, 0)
        view = HintedTreeView()
        swin.add(view)

        def icon_cdf(column, cell, model, iter_, *args):
            item = model[iter_][0]
            if item.error:
                cell.set_property('stock-id', Gtk.STOCK_DIALOG_ERROR)
            else:
                cell.set_property('stock-id', None)

        column = Gtk.TreeViewColumn()
        column.set_sizing(Gtk.TreeViewColumnSizing.AUTOSIZE)
        icon_render = Gtk.CellRendererPixbuf()
        column.pack_start(icon_render, True)
        column.set_cell_data_func(icon_render, icon_cdf)
        view.append_column(column)

        def track_cdf(column, cell, model, iter_, *args):
            item = model[iter_][0]
            cell.set_property('text', item.title)
            cell.set_sensitive(model[iter_][1])

        column = Gtk.TreeViewColumn(_("Track"))
        column.set_expand(True)
        column.set_sizing(Gtk.TreeViewColumnSizing.AUTOSIZE)
        track_render = Gtk.CellRendererText()
        track_render.set_property('ellipsize', Pango.EllipsizeMode.END)
        column.pack_start(track_render, True)
        column.set_cell_data_func(track_render, track_cdf)
        view.append_column(column)

        def progress_cdf(column, cell, model, iter_, *args):
            item = model[iter_][0]
            cell.set_property('value', int(item.progress * 100))
            cell.set_sensitive(model[iter_][1])

        column = Gtk.TreeViewColumn(_("Progress"))
        column.set_sizing(Gtk.TreeViewColumnSizing.AUTOSIZE)
        progress_render = Gtk.CellRendererProgress()
        column.pack_start(progress_render, True)
        column.set_cell_data_func(progress_render, progress_cdf)
        view.append_column(column)

        def gain_cdf(column, cell, model, iter_, *args):
            item = model[iter_][0]
            if item.gain is None or not item.done:
                cell.set_property('text', "-")
            else:
                cell.set_property('text', "%.2f db" % item.gain)
            cell.set_sensitive(model[iter_][1])

        column = Gtk.TreeViewColumn(_("Gain"))
        column.set_sizing(Gtk.TreeViewColumnSizing.AUTOSIZE)
        gain_renderer = Gtk.CellRendererText()
        column.pack_start(gain_renderer, True)
        column.set_cell_data_func(gain_renderer, gain_cdf)
        view.append_column(column)

        def peak_cdf(column, cell, model, iter_, *args):
            item = model[iter_][0]
            if item.gain is None or not item.done:
                cell.set_property('text', "-")
            else:
                cell.set_property('text', "%.2f" % item.peak)
            cell.set_sensitive(model[iter_][1])

        column = Gtk.TreeViewColumn(_("Peak"))
        column.set_sizing(Gtk.TreeViewColumnSizing.AUTOSIZE)
        peak_renderer = Gtk.CellRendererText()
        column.pack_start(peak_renderer, True)
        column.set_cell_data_func(peak_renderer, peak_cdf)
        view.append_column(column)

        self.create_pipelines()
        self._timeout = None
        self._sigs = {}
        self._done = []

        self.__fill_view(view, albums)
        num_to_process = sum(int(rga.should_process) for rga in self._todo)
        template = ngettext("There is <b>%d</b> album to update (of %d)",
                            "There are <b>%d</b> albums to update (of %d)",
                            num_to_process)
        info.set_markup(template % (num_to_process, len(self._todo)))
        self.connect("destroy", self.__destroy)
        self.connect('response', self.__response)

    def create_pipelines(self):
        # create as many pipelines as threads
        self.pipes = [ReplayGainPipeline() for _ in xrange(get_num_threads())]

    def __fill_view(self, view, albums):
        self._todo = [RGAlbum.from_songs(a, self.process_mode) for a in albums]
        self._count = len(self._todo)
        self.model = model = Gtk.TreeStore(object, bool)
        insert = model.insert
        for album in reversed(self._todo):
            enabled = album.should_process
            base = insert(None, 0, row=[album, enabled])
            for song in reversed(album.songs):
                insert(base, 0, row=[song, enabled])
        view.set_model(model)

        if len(self._todo) == 1:
            view.expand_all()

    def start_analysis(self):
        self._timeout = GLib.idle_add(self.__request_update)

        # fill the pipelines
        for p in self.pipes:
            if not self._todo:
                break
            self._sigs[p] = [
                p.connect("done", self.__done),
                p.connect("update", self.__update),
            ]
            album = self.get_next_album()
            if not album:
                return
            p.start(album)

    def get_next_album(self):
        next_album = None
        while not next_album:
            if not self._todo:
                print_d("No more albums to process")
                return None
            next_album = self._todo.pop(0)
            if not next_album.should_process:
                print_d("%s needs no processing" % next_album.title)
                self._done.append(next_album)
                self.__update_view_for(next_album)
                next_album = None
        return next_album

    def __response(self, win, response):
        if response == Gtk.ResponseType.CANCEL:
            self.destroy()
        elif response == Gtk.ResponseType.OK:
            for album in self._done:
                album.write()
            self.destroy()

    def __destroy(self, *args):
        # shut down any active processing and clean up resources, timeouts
        if self._timeout:
            GLib.source_remove(self._timeout)
        for p in self.pipes:
            if p in self._sigs:
                for s in self._sigs.get(p, []):
                    p.disconnect(s)
            p.quit()

    def __update(self, pipeline, album, song):
        for row in self.model:
            row_album = row[0]
            if row_album is album:
                self.model.row_changed(row.path, row.iter)
                for child in row.iterchildren():
                    row_song = child[0]
                    if row_song is song:
                        self.model.row_changed(child.path, child.iter)
                        break
                break

    def __done(self, pipeline, album):
        self._done.append(album)
        if self._todo:
            pipeline.start(self._todo.pop(0))
        self.__update_view_for(album)

    def __update_view_for(self, album):
        for row in self.model:
            row_album = row[0]
            if row_album is album:
                self.model.row_changed(row.path, row.iter)
                break

    def __request_update(self):
        GLib.source_remove(self._timeout)
        self._timeout = None
        # all done, stop
        if len(self._done) < self._count:
            for p in self.pipes:
                p.request_update()
            self._timeout = GLib.timeout_add(400, self.__request_update)
        return False


class ReplayGain(SongsMenuPlugin, PluginConfigMixin):
    PLUGIN_ID = 'ReplayGain'
    PLUGIN_NAME = _('Replay Gain')
    PLUGIN_DESC = _('Analyzes and updates ReplayGain information, '
                    'using GStreamer. Results are grouped by album.')
    PLUGIN_ICON = Gtk.STOCK_MEDIA_PLAY
    # A little arbitrary, but seemed to match the history best.
    PLUGIN_VERSION = '2.1'
    CONFIG_SECTION = 'replaygain'

    def plugin_albums(self, albums):
        mode = self.config_get("process_if", UpdateMode.ALWAYS)
        win = RGDialog(albums, parent=self.plugin_window, process_mode=mode)
        win.show_all()
        win.start_analysis()

        # plugin_done checks for metadata changes and opens the write dialog
        win.connect("destroy", self.__plugin_done)

    def __plugin_done(self, win):
        self.plugin_finish()

    @classmethod
    def PluginPreferences(cls, parent):
        vb = Gtk.VBox(spacing=12)

        # Server settings Frame
        frame = Gtk.Frame(label=_("<b>Existing Tags</b>"))
        frame.set_shadow_type(Gtk.ShadowType.NONE)
        frame.get_label_widget().set_use_markup(True)
        frame_align = Gtk.Alignment.new(0, 0, 1, 1)
        frame_align.set_padding(6, 6, 12, 12)
        frame.add(frame_align)

        # Tabulate all settings for neatness
        table = Gtk.Table(n_rows=1, n_columns=2)
        table.set_col_spacings(6)
        table.set_row_spacings(6)
        rows = []

        def process_option_changed(combo):
            #xcode = combo.get_child().get_text()
            model = combo.get_model()
            lbl, value = model[combo.get_active()]
            cls.config_set("process_if", value)

        def create_model():
            model = Gtk.ListStore(str, str)
            model.append([_("<b>always</b>"), UpdateMode.ALWAYS])
            model.append([_("if <b>any</b> RG tags are missing"),
                          UpdateMode.ANY_MISSING])
            model.append([_("if <b>album</b> RG tags are missing"),
                          UpdateMode.ALBUM_MISSING])
            return model

        def set_active(value):
            for i, item in enumerate(model):
                if value == item[1]:
                    combo.set_active(i)

        model = create_model()
        combo = Gtk.ComboBox(model=model)
        set_active(cls.config_get("process_if", UpdateMode.ALWAYS))
        renderer = Gtk.CellRendererText()
        combo.connect('changed', process_option_changed)
        combo.pack_start(renderer, True)
        combo.add_attribute(renderer, "markup", 0)

        rows.append((_("_Process albums:"), combo))

        for (row, (label_text, entry)) in enumerate(rows):
            label = Gtk.Label(label=label_text)
            label.set_alignment(0.0, 0.5)
            label.set_use_underline(True)
            label.set_mnemonic_widget(entry)
            table.attach(label, 0, 1, row, row + 1,
                         xoptions=Gtk.AttachOptions.FILL)
            table.attach(entry, 1, 2, row, row + 1)

        frame_align.add(table)
        vb.pack_start(frame, True, True, 0)
        return vb


if not Gst.Registry.get().find_plugin("replaygain"):
    __all__ = []
    del ReplayGain
    raise ImportError("GStreamer replaygain plugin not found")
