import AcraNetwork.Pcap as pcap
import AcraNetwork.iNetX as inetx
import AcraNetwork.MPEGTS as mpegts
import AcraNetwork.SimpleEthernet as eth
from AcraNetwork.MPEGTS import H264
import datetime

# This script shows how to parse either a pcap file or a TS file into the constituent
# NALS and finds the unique STANAG SEI User data with the timestamp

def pcap_to_ts(pcapfile,ts_file,udp_port=8010):
    '''
    Convert a pcap file to a TS file by extracting all data from a specified port
    :param mpegfile: str
    :param ts_file: str
    :param udp_port: int
    :return:
    '''

    mpegpcap = pcap.Pcap(pcapfile, mode='r')
    mpegpcap.readGlobalHeader()

    ts = open(ts_file, mode='wb')
    mpeghex = ""
    rec_count = 0

    while not ts.closed:
        try:
            rec = mpegpcap.readAPacket()
        except:
            ts.close()
        else:
            try:
                e = eth.Ethernet()
                e.unpack(rec.packet)
                i = eth.IP()
                i.unpack(e.payload)
                u = eth.UDP()
                u.unpack(i.payload)
                if u.dstport == udp_port:
                    rec_count += 1
                    inet = inetx.iNetX()
                    inet.unpack(u.payload)
                    mpegtspackets = mpegts.MPEGTS()
                    mpegtspackets.unpack(inet.payload)
                    for packet in mpegtspackets.blocks:
                        mpeghex += packet.payload
                    ts.write(inet.payload)
            except:
                continue


def parse_ts_file(tsfile):

    ts_file = open(tsfile,mode='rb')
    h264_data = H264()
    h264_data.unpack(ts_file.read())

    nal_counts = {}
    timestamp_count = 0
    for nal in h264_data.nals:

        if not nal.type in nal_counts:
            nal_counts[nal.type] = 1
        else:
            nal_counts[nal.type] += 1

        if nal.type == mpegts.NAL_TYPES["SEI"]:
            if nal.sei.unregdata:
                print "Timestamp={} byte offset={} count ={}".format(datetime.datetime.strftime(nal.sei.time,"%d %b %Y %H:%M:%S.%f"),nal.offset,timestamp_count)
                timestamp_count += 1

        else:
            pass

    print "\n----- SUMMARY -----"
    for type in nal_counts:
        print "{} {} NALs in input".format(nal_counts[type],mpegts.NAL_TYPES_INV[type])
    print "{} STANAG Timestamps".format(timestamp_count)

def main():
    # Read in a TS file and print out some useful information
    parse_ts_file("../test/stanag_sample.ts")


if __name__ == "__main__":
    main()