#!/usr/bin/python

import sys

import gi
gi.require_version( 'Gst', '1.0' )
from gi.repository import Gst, GLib, GObject
GObject.threads_init()
GLib.threads_init()
Gst.init(None)

import threading

print_lock = threading.Lock()

def output( string ):
    global print_lock
    with print_lock:
        print string

class TestPipeline:
    def __init__( self, filename ):
        self.last_pts = {}
        self.last_dts = {}

        self.pipeline = Gst.Pipeline( 'PTS Check' )
        bus = self.pipeline.get_bus()
        bus.add_signal_watch()
        bus.connect( 'message::error', self.__pipeline_error )
        bus.connect( 'message::warning', self.__pipeline_warning )
        bus.connect( 'message::eos', self.__pipeline_eos )
        bus.connect( 'message', self.__pipeline_message )

        # filesrc ! mpegpsdemux ! mpegvideoparse ! vaapidecode ! vaapipostproc deinterlace-method=motion-adaptive ! vaapisink
        self.filesrc = Gst.ElementFactory.make( 'filesrc' )
        self.filesrc.set_property( 'location', filename )
        self.addElement( self.filesrc )

        self.mpegpsdemux = Gst.ElementFactory.make( 'mpegpsdemux' )
        self.addElement( self.mpegpsdemux )

        self.mpegvideoparse = Gst.ElementFactory.make( 'mpegvideoparse' )
        self.addElement( self.mpegvideoparse )
        self.addPtsProbe( 'mpegvideoparse output', self.mpegvideoparse.get_static_pad( 'src' ) )

        self.vaapidecode = Gst.ElementFactory.make( 'vaapidecode' )
        self.addElement( self.vaapidecode )
        self.addPtsProbe( 'vaapidecode output', self.vaapidecode.get_static_pad( 'src' ) )

        self.vaapipostproc = Gst.ElementFactory.make( 'vaapipostproc' )
        self.vaapipostproc.set_property( 'deinterlace-method', 'motion-adaptive' )
        self.addElement( self.vaapipostproc )
        self.addPtsProbe( 'vaapipostproc output', self.vaapipostproc.get_static_pad( 'src' ) )

        self.vaapisink = Gst.ElementFactory.make( 'vaapisink' )
        self.addElement( self.vaapisink )

        self.filesrc.link( self.mpegpsdemux )
        self.mpegpsdemux.connect( 'pad-added', self.__mpegpsdemux_pad_added )
        self.mpegvideoparse.link( self.vaapidecode )
        self.vaapidecode.link( self.vaapipostproc )
        self.vaapipostproc.link( self.vaapisink )

    def __mpegpsdemux_pad_added( self, element, pad ):
        output( 'New pad %r %r' % (element, pad,) )
        pad.link( self.mpegvideoparse.get_static_pad( 'sink' ) )

    def __pipeline_error( self, bus, message ):
        (error, debug) = message.parse_error()
        output( 'ERROR from %r: %r debug %r' % (message.src, error, debug,) )

    def __pipeline_warning( self, bus, message ):
        (warning, debug) = message.parse_warning()
        output( 'WARNING from %r: %r debug %r' % (message.src, warning, debug,) )

    def __pipeline_message( self, bus, message ):
        contents = None
        structure = message.get_structure()
        name = None
        if structure is not None:
            name = structure.get_name()

        if message.type == Gst.MessageType.ASYNC_DONE:
            contents = message.parse_async_done()

        elif message.type == Gst.MessageType.HAVE_CONTEXT:
            contents = message.parse_have_context()

        elif message.type == Gst.MessageType.NEW_CLOCK:
            contents = message.parse_new_clock()

        elif message.type == Gst.MessageType.STATE_CHANGED:
            contents = message.parse_state_changed()

        elif message.type == Gst.MessageType.STREAM_STATUS:
            contents = message.parse_stream_status()

        elif message.type == Gst.MessageType.TAG:
            contents = message.parse_tag()

        elif message.type == Gst.MessageType.WARNING:
            contents = message.parse_warning()

        elif message.type == Gst.MessageType.ERROR:
            contents = message.parse_error()

        elif message.type == Gst.MessageType.QOS:
            # Note the lack of () around message.parse is deliberate
            contents = 'live_stream %r running_time %r stream_time %r timestamp %r duration %r ' % message.parse_qos()
            contents += 'jitter %r proportion %r quality %r ' % message.parse_qos_values()
            contents += 'format %r processed %r dropped %r' % message.parse_qos_stats()

        output( 'Async Message %s contents %r from %r' % (name, contents, message.src,) )

    def quit( self ):
        output( 'Quitting' )
        global main_loop
        main_loop.quit()

    def __pipeline_eos( self, bus, message ):
        output( 'EOS' )
        GLib.timeout_add_seconds( 10, self.quit )

    def __pts_probe( self, pad, info, name ):
        invalid_time = 18446744073709551615
        buf = info.get_buffer()

        old_pts = self.last_pts[ name ]
        self.last_pts[ name ] = buf.pts
        if old_pts == invalid_time:
            pts_diff = 0
        else:
            pts_diff = buf.pts - old_pts

        old_dts = self.last_dts[ name ]
        self.last_dts[ name ] = buf.dts
        if old_dts == invalid_time:
            dts_diff = 0
        else:
            dts_diff = buf.dts - old_dts

        if buf.pts == invalid_time and buf.dts == invalid_time:
            output( 'pts probe %30s no timestamps' % (name,) )
        elif buf.pts == invalid_time:
            output( 'pts probe %30s          %15s     %15s dts_diff %15d dts %15d'  % (name,
                                                                                       '', '',
                                                                                       dts_diff, buf.dts) )
        elif buf.dts == invalid_time:
            output( 'pts probe %30s pts_diff %15d pts %15d'  % (name,
                                                                buf.pts - old_pts, buf.pts,) )
        else:
            output( 'pts probe %30s pts_diff %15d pts %15d dts_diff %15d dts %15d'  % (name,
                                                                                       pts_diff, buf.pts,
                                                                                       dts_diff, buf.dts) )
        return Gst.PadProbeReturn.OK

    def addElement( self, element ):
        self.pipeline.add( element )
        element.sync_state_with_parent()

    def addPtsProbe( self, name, pad ):
        self.last_pts[ name ] = 0
        self.last_dts[ name ] = 0
        pad.add_probe( Gst.PadProbeType.BUFFER, self.__pts_probe, name )

    def play( self ):
        self.pipeline.set_state( Gst.State.PLAYING )

def main( argv ):
    global main_loop
    main_loop = GLib.MainLoop()

    if len(sys.argv) != 2:
        print 'Usage: %s <file>' % (argv[0],)
        return 1

    pipeline = TestPipeline( argv[1] )

    pipeline.play()

    main_loop.run()
    return 0

if __name__ == "__main__":
    sys.exit( main( sys.argv ) )
