import argparse,hashlib,json,time
from pathlib import Path
import gi
gi.require_version('Gst','1.0');gi.require_version('GstRtp','1.0')
from gi.repository import Gst,GstRtp,GLib
Gst.init(None)
p=argparse.ArgumentParser();p.add_argument('uri');p.add_argument('seconds',type=int);p.add_argument('output');a=p.parse_args()
loop=GLib.MainLoop();pipe=Gst.Pipeline.new('bounded-receiver');src=Gst.ElementFactory.make('rtspsrc');sink=Gst.ElementFactory.make('fakesink')
src.set_property('location',a.uri);src.set_property('protocols',4);src.set_property('latency',50)
sink.set_property('sync',False);pipe.add(src);pipe.add(sink)
report=dict(packets=0,clock_rate=None,sdp=False,frames=[],sender_reports=[],errors=[])
manager=None;current=None;start=time.monotonic();rtcp_hooked=False
def flush_frame():
    global current
    if current is not None and len(report['frames'])<180:
        row=dict(current);row['digest']=row['digest'].hexdigest();report['frames'].append(row)
    current=None
def rtp_probe(pad,info):
    global current
    buf=info.get_buffer();ok,packet=GstRtp.RTPBuffer.map(buf,Gst.MapFlags.READ)
    if not ok:return Gst.PadProbeReturn.OK
    ts,seq,ssrc=packet.get_timestamp(),packet.get_seq(),packet.get_ssrc();packet.unmap()
    report['packets']+=1
    if current is None or ts!=current['rtp']:
        flush_frame();current=dict(rtp=ts,ssrc=ssrc,first_seq=seq,last_seq=seq,digest=hashlib.sha256(),elapsed=time.monotonic()-start)
    current['last_seq']=seq
    ok,m=buf.map(Gst.MapFlags.READ)
    if ok:current['digest'].update(m.data);buf.unmap(m)
    if len(report['frames'])>=180:GLib.idle_add(loop.quit)
    return Gst.PadProbeReturn.OK
def rtcp_probe(pad,info):
    b=GstRtp.RTCPBuffer()
    if GstRtp.RTCPBuffer.map(info.get_buffer(),Gst.MapFlags.READ,b):
        try:
            packet=GstRtp.RTCPPacket();valid=b.get_first_packet(packet)
            while valid:
                if packet.get_type()==GstRtp.RTCPType.SR and len(report['sender_reports'])<12:
                    ssrc,ntp,rtp,packets,octets=packet.sr_get_sender_info()
                    report['sender_reports'].append(dict(ssrc=ssrc,ntp=str(ntp),rtp=rtp,packets=packets))
                valid=packet.move_to_next()
        finally:b.unmap()
    return Gst.PadProbeReturn.OK
def new_manager(_src,value):
    global manager
    manager=value
def pad_added(_src,pad):
    global rtcp_hooked
    caps=pad.get_current_caps().get_structure(0)
    if caps.get_value('media')!='video':return
    report['clock_rate']=caps.get_value('clock-rate');report['encoding']=caps.get_value('encoding-name')
    pad.add_probe(Gst.PadProbeType.BUFFER,rtp_probe);pad.link(sink.get_static_pad('sink'))
    if manager is not None:
        rp=manager.get_static_pad('recv_rtcp_sink_0')
        if rp is not None:rp.add_probe(Gst.PadProbeType.BUFFER,rtcp_probe);rtcp_hooked=True
def message(_bus,msg):
    if msg.type==Gst.MessageType.ERROR:
        error,detail=msg.parse_error();report['errors'].append(str(error));report['error_detail']=detail;loop.quit()
    elif msg.type==Gst.MessageType.EOS:report['eos']=True;loop.quit()
src.connect('new-manager',new_manager);src.connect('pad-added',pad_added);src.connect('on-sdp',lambda *_:report.update(sdp=True))
bus=pipe.get_bus();bus.add_signal_watch();bus.connect('message',message)
GLib.timeout_add_seconds(a.seconds,lambda:(loop.quit(),False)[1]);pipe.set_state(Gst.State.PLAYING)
try:loop.run()
finally:
    pipe.set_state(Gst.State.NULL);flush_frame();report.update(elapsed=time.monotonic()-start,rtcp_hooked=rtcp_hooked)
    Path(a.output).write_text(json.dumps(report,indent=2))
    print(json.dumps({k:v for k,v in report.items() if k not in ('frames','sender_reports')}))