#!/usr/bin/env python

"""
A wrapper around playbin for playing DASH streams.

This script allows playback of DASH streams and provides basic
trickplay controls such as seek, fast-forward and fast-rewind.

Author: Alex Ashley <alex.ashley@youview.com>
"""

import argparse, sys

import gi
gi.require_version('Gst', '1.0')
from gi.repository import GObject, Gst, Gtk

# Needed for window.get_xid(), xvimagesink.set_window_handle(), respectively:
from gi.repository import GdkX11, GstVideo

GObject.threads_init()
Gst.init(None)

class VideoWidget(Gtk.DrawingArea):
    """
    Extend gtk.DrawingArea to create our own video widget.
    http://www.majorsilence.com/pygtk_audio_and_video_playback_gstreamer
    """
    def __init__(self):
        Gtk.DrawingArea.__init__(self)
        self.imagesink = None
        self.set_double_buffered(False)

    def do_expose_event(self, event):
        if self.imagesink:
            self.imagesink.expose()
            return False
        else:
            return True

    def set_sink(self, sink, win_id):
        self.imagesink = sink
        self.imagesink.set_window_handle(win_id)

class Player(object):
    TRICKPLAY_RATES=[2.0, 6.0, 12.0, 32.0, 64.0]
    def __init__(self, args):
        self.current_position=None
        self.duration = None
        self.mode = None
        self.seek_range = None
        self.is_playing = False
        self.is_ready = False
        self.is_seeking = False
        self.playback_rate = 1.0
        self.trickplay_amount = args.trickplay
        mpd_url = args.url
        self.window = Gtk.Window()
        self.window.connect('destroy', self.quit)
        self.window.set_default_size(800, 450)
        self.window.set_title(mpd_url)

        vbox = Gtk.VBox(False, 0) 
        self.video_window = VideoWidget()
        self.video_window.set_size_request(800,450)
        vbox.pack_start(self.video_window, False, True, 0) 
        self.scale_range = 100.0
        self.adjustment = Gtk.Adjustment(0.0, 0.00, self.scale_range, 0.1, 1.0, 1.0)
        self.timeline = Gtk.Scale(orientation=Gtk.Orientation.HORIZONTAL, adjustment=self.adjustment)
        self.timeline.set_hexpand(True)
        self.timeline.set_draw_value(False)
        self.value_changed_signal = self.timeline.connect("value-changed", self.timeline_change)
        vbox.pack_start(self.timeline, False, True, 0) 
        hbox = Gtk.Box(False, 0, orientation=Gtk.Orientation.HORIZONTAL, spacing=6)

        playButtonImage = Gtk.Image()
        playButtonImage.set_from_stock("gtk-media-play", Gtk.IconSize.BUTTON)
        self.play_button = Gtk.Button.new()
        self.play_button.add(playButtonImage)
        self.play_button.connect("clicked", self.on_play_clicked)
        hbox.pack_start(self.play_button, expand=False, fill=True, padding=2)

        prevButtonImage = Gtk.Image()
        prevButtonImage.set_from_stock("gtk-media-previous", Gtk.IconSize.BUTTON)
        self.previous_button = Gtk.Button.new()
        self.previous_button.add(prevButtonImage)
        self.previous_button.connect("clicked", self.on_previous_clicked)
        hbox.pack_start(self.previous_button, False, True, 0)

        img = Gtk.Image()
        img.set_from_stock("gtk-media-rewind", Gtk.IconSize.BUTTON)
        self.rwnd_button = Gtk.Button.new()
        self.rwnd_button.add(img)
        self.rwnd_button.connect("clicked", self.on_rwnd_clicked)
        hbox.pack_start(self.rwnd_button, expand=False, fill=True, padding=2)

        img = Gtk.Image()
        img.set_from_stock("gtk-media-pause", Gtk.IconSize.BUTTON)
        self.pause_button = Gtk.Button.new()
        self.pause_button.add(img)
        self.pause_button.connect("clicked", self.on_pause_clicked)
        hbox.pack_start(self.pause_button, False, True, 0)

        img = Gtk.Image()
        img.set_from_stock("gtk-media-forward", Gtk.IconSize.BUTTON)
        self.fwd_button = Gtk.Button.new()
        self.fwd_button.add(img)
        self.fwd_button.connect("clicked", self.on_fwd_clicked)
        hbox.pack_start(self.fwd_button, False, True, 0)

        img = Gtk.Image()
        img.set_from_stock("gtk-media-next", Gtk.IconSize.BUTTON)
        self.next_button = Gtk.Button.new()
        self.next_button.add(img)
        self.next_button.connect("clicked", self.on_next_clicked)
        hbox.pack_start(self.next_button, False, True, 0)

        img = Gtk.Image()
        img.set_from_stock("gtk-media-stop", Gtk.IconSize.BUTTON)
        self.stop_button = Gtk.Button.new()
        self.stop_button.add(img)
        self.stop_button.connect("clicked", self.on_stop_clicked)
        hbox.pack_start(self.stop_button, False, True, 0)

        img = Gtk.Image()
        img.set_from_stock("gtk-close", Gtk.IconSize.BUTTON)
        self.exit_button = Gtk.Button.new()
        self.exit_button.add(img)
        self.exit_button.connect("clicked", self.quit)
        hbox.pack_start(self.exit_button, False, True, 0)

        self.time_label = Gtk.Label("00:00:00 / 00:00:00")
        hbox.pack_end(self.time_label, False, True, padding=20)

        self.rate_label = Gtk.Label("x  1")
        hbox.pack_end(self.rate_label, False, True, padding=10)

        vbox.pack_start(hbox, False, True, 0) 
        self.window.add(vbox)

        if args.debug:
            Gst.debug_set_active(True)
            Gst.debug_set_default_threshold(2)
            Gst.debug_set_threshold_for_name('adaptivedemux',5)
            Gst.debug_set_threshold_for_name('dashdemux',5)
            Gst.debug_set_threshold_for_name('qtdemux',4)
            Gst.debug_set_threshold_for_name('cencdec',3)

        # Create GStreamer pipeline
        self.pipeline = Gst.Pipeline()

        self.time_format = Gst.Format.TIME

        # Create bus to get events from GStreamer pipeline
        self.bus = self.pipeline.get_bus()
        self.bus.add_signal_watch()
        self.bus.connect('message::eos', self.on_eos)
        self.bus.connect('message::error', self.on_error)
        self.bus.connect('message::state-changed', self.on_state_change)

        # This is needed to make the video output in our DrawingArea:
        self.bus.enable_sync_message_emission()
        self.bus.connect('sync-message::element', self.on_sync_message)

        # Create GStreamer elements
        self.playbin = Gst.ElementFactory.make('playbin', None)
        if not self.playbin:
            print('Problem with gstreamer installation: playbin element not found.')
            sys.exit(2)
        self.pipeline.add(self.playbin)
        self.playbin.set_property('uri', mpd_url)

    def run(self):
        self.window.show_all()
        # You need to get the XID after window.show_all().  You shouldn't get it
        # in the on_sync_message() handler because threading issues will cause
        # segfaults there.
        self.xid = self.video_window.get_property('window').get_xid()
        self.pipeline.set_state(Gst.State.PLAYING)
        self.update_timer = GObject.timeout_add(1000, self.update_time_label)
        Gtk.main()

    def seek(self, delta=None, position=None):
        if self.is_playing and (delta is not None or position is not None):
            new_pos = position
            self.update_seek_range()
            if delta:
                self.update_current_position()
                new_pos = self.current_position+delta
                if new_pos < self.seek_range[0]:
                    new_pos = self.seek_range[0]
                elif new_pos > self.seek_range[1]:
                    new_pos = self.seek_range[1]
            self.is_seeking = True
            self.current_position = new_pos
            if not self.pipeline.seek_simple(
                    Gst.Format.TIME,        
                    Gst.SeekFlags.FLUSH | Gst.SeekFlags.KEY_UNIT,
                    new_pos):
                self.is_seeking = False
            self.update_time_label()

    def set_playback_rate(self,rate):
        self.playback_rate=rate
        self.rate_label.set_text('x %1.0f'%self.playback_rate)
        if not self.is_playing:
            return
        current_position = self.pipeline.query_position(self.time_format)
        if current_position:
            start_type = Gst.SeekType.SET
            start_pos = current_position[1] if rate > 0 else 0
            stop_type = Gst.SeekType.SET if rate < 0 else Gst.SeekType.NONE
            stop_pos = current_position[1] if rate < 0 else -1
        else:
            start_pos = stop_pos = -1
            start_type = stop_type = Gst.SeekType.NONE
        self.playbin.set_property('mute',rate!=1)
        self.pipeline.seek(rate=self.playback_rate, \
                           format=self.time_format, \
                           flags=Gst.SeekFlags.FLUSH, \
                           start_type=start_type, \
                           start=start_pos, \
                           stop_type=stop_type, \
                           stop=stop_pos)

    def timeline_change(self,event):
        self.update_seek_range()
        if self.seek_range is not None and self.current_position:
            pos = self.timeline.get_value()
            r = self.seek_range[1] - self.seek_range[0]
            pos = (r * pos / self.scale_range) + self.seek_range[0]
            if pos>self.seek_range[1]:
                pos=self.seek_range[1]
            if abs(pos-self.current_position)>Gst.SECOND:
                self.seek(position=pos)

    def quit(self, window):
        self.pipeline.set_state(Gst.State.NULL)
        Gtk.main_quit()

    def on_sync_message(self, bus, msg):
        #print(msg.get_structure().get_name())
        if msg.get_structure().get_name() == 'prepare-window-handle':
            #print('prepare-window-handle')
            self.video_window.set_sink(msg.src, self.xid)
            #msg.src.set_window_handle(self.xid)

    def on_state_change(self, bus, msg):
        if msg.src==self.playbin:
            #print('state-change',msg.src)
            old_state, new_state, pending = msg.parse_state_changed()
            #print(old_state,new_state,pending)
            if not self.is_ready and new_state>=Gst.State.PAUSED:
                self.is_ready=True
                print('gstreamer pipeline ready')
                elt = self.find_element_by_name('dashdemux0')
                if elt and self.trickplay_amount:
                    elt.set_property("trickplay-time", float(self.trickplay_amount))
            elif self.is_ready:
                self.is_playing = new_state==Gst.State.PLAYING
                if self.is_playing and self.is_seeking:
                    self.is_seeking=False
                    print('Seek complete')
                if new_state < Gst.State.READY:
                    self.is_ready=False

    def on_eos(self, bus, msg):
        if self.seek_range:
            print('on_eos(): seeking to start of video')
            self.pipeline.seek_simple(
                Gst.Format.TIME,
                Gst.SeekFlags.FLUSH | Gst.SeekFlags.KEY_UNIT,
                self.seek_range[0]
            )

    def on_error(self, bus, msg):
        print('on_error():', msg.parse_error())

    def update_seek_range(self):
        if self.seek_range is None or self.mode!='vod':
            query = Gst.Query.new_seeking(Gst.Format.TIME)
            if self.pipeline.query(query):
                (fmt, self.seek_enabled, start, end) = query.parse_seeking()
                self.seek_range = (start,end)
                if start>0:
                    self.mode='live'

    def update_current_position(self):
        if not self.is_seeking:
            # Query position during a seek returns 0
            self.current_position = self.pipeline.query_position(self.time_format)[1]
        #print(self.current_position)

    def update_time_label(self): 
        """
        Update the time_label to display the current location
        in the media file as well as update the seek bar
        """ 
        if not self.is_playing:
            #print('not playing')
            return True
        self.update_seek_range()
        if self.duration is None and self.mode!='live':
            try:
                self.duration = self.pipeline.query_duration(self.time_format)
                if self.duration and not self.duration[1]:
                    self.duration = None
                else:
                    self.duration_str = self.convert_time(self.duration[1])
                    if self.seek_range and self.seek_range[0]==0:
                        self.mode='vod'
                #print('Duration',self.duration)
            except Exception,e:
                print(str(e))
                self.duration = None
                if self.seek_range and self.seek_range[0]>0:
                    self.mode='live'
        if self.mode=='live' or (self.duration is not None and self.duration[1]):
            self.update_current_position()
        if self.seek_range is not None and self.current_position is not None:
            current_position_formated = self.convert_time(self.current_position)
            if self.mode!='live' and self.duration is not None and self.duration[1]:
                self.time_label.set_text(current_position_formated + "/" + self.duration_str)
            else:
                start_pos = self.convert_time(self.seek_range[0])
                end_pos = self.convert_time(self.seek_range[1])
                self.time_label.set_text('[' + start_pos + '] ' + current_position_formated + ' [' + end_pos + ']')
            r = self.seek_range[1] - self.seek_range[0]
            if r>0:
                pos = self.scale_range * (self.current_position-self.seek_range[0]) / r
                #print self.mode,self.seek_range,self.current_position,r,pos
                if pos<0:
                    pos=0
            # don't trigger the value-change handler when updating
            # the slider
            self.timeline.handler_block(self.value_changed_signal)
            self.timeline.set_value(pos)
            self.timeline.handler_unblock(self.value_changed_signal)
        return True

    def convert_time(self, ns=None): 
        if ns is None: 
            return None 
        seconds = ns / Gst.SECOND
        minutes = seconds // 60
        hours = minutes // 60
        minutes = minutes % 60
        seconds = seconds % 60
        return '%02d:%02d:%02d'%(hours,minutes,seconds)

    def find_element_by_name(self, name, gstbin=None):
        if gstbin is None:
            gstbin=self.playbin
        elt = gstbin.get_by_name(name)
        if elt is None:
            for e in gstbin.iterate_elements():
                #print e.get_name()
                if isinstance(e,Gst.Bin):
                    a = self.find_element_by_name(name,e)
                    if a:
                        return a
        return elt
            
    def on_pause_clicked(self, widget):
        self.pipeline.set_state(Gst.State.PAUSED)

    def on_play_clicked(self, widget):
        self.pipeline.set_state(Gst.State.PLAYING)
        if self.playback_rate!=1.0:
            self.set_playback_rate(1.0)

    def on_stop_clicked(self, widget):
        self.pipeline.set_state(Gst.State.NULL)

    def on_next_clicked(self, widget):
        self.seek(delta=20*Gst.SECOND)

    def on_previous_clicked(self, widget):
        self.seek(delta=(-20*Gst.SECOND))

    def on_rwnd_clicked(self,widget):
        for rate in self.TRICKPLAY_RATES:
            if self.playback_rate > -rate:
                self.set_playback_rate(-rate)
                return

    def on_fwd_clicked(self,widget):
        for rate in self.TRICKPLAY_RATES:
            if self.playback_rate<rate:
                self.set_playback_rate(rate)
                return

def usage():
    print("Usage:")
    print("    %s <MPD URL>"%(sys.argv[0]))

parser = argparse.ArgumentParser(description='DASH player')
parser.add_argument('-d', '--debug', action="store_true")
parser.add_argument('-t', '--trickplay', action='store', default=None, \
                    help='amount of a fragment to download when in trick play mode (in seconds)')
parser.add_argument('url', help='MPD URL')

print("Starting player...")
p = Player(parser.parse_args())
p.run()
