#! /usr/bin/env python3
# -*- coding: utf-8 -*-

#
# @RDE_HEADER@
# $Id$
# $HeadURL$
# $Author$
# Created by Christian Fröbel, NLE, 2019-01-29
#

import ctypes
import socket
import sys
import textwrap
import time
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from collections import namedtuple
from contextlib import ExitStack, contextmanager
from queue import Empty, Queue
from threading import Thread

_G711_PAYLOAD_TYPE = 0x08


@contextmanager
def autoShutdown(target):
    try:
        yield target
    finally:
        target.initShutdown()

@contextmanager
def autoJoin(thread):
    try:
        yield thread
    finally:
        thread.join()

@contextmanager
def startThreadContext(target, thread):
    with ExitStack() as stack:
        thread.start()
        stack.enter_context(autoJoin(thread))
        stack.enter_context(autoShutdown(target))
        yield


class _RtpHeaderStruct(ctypes.BigEndianStructure):
    _pack_ = 1
    _fields_ = [
        ('version',        ctypes.c_uint, 2),
        ('padding',        ctypes.c_uint, 1),
        ('extension',      ctypes.c_uint, 1),
        ('csrcCount',      ctypes.c_uint, 4),
        ('marker',         ctypes.c_uint, 1),
        ('payloadType',    ctypes.c_uint, 7),
        ('sequenceNumber', ctypes.c_uint, 16),
        ('timestamp',      ctypes.c_uint, 32),
        ('ssrc',           ctypes.c_uint, 32)
    ]

class _RtpHeader(object):
    def __init__(self, payloadType, sequenceNumber, timestamp, ssrc):
        self._buf = bytearray(ctypes.sizeof(_RtpHeaderStruct))
        self._header = _RtpHeaderStruct.from_buffer(self._buf)
        self._setFixedValues()
        self._header.payloadType    = payloadType
        self._header.sequenceNumber = sequenceNumber
        self._header.timestamp      = timestamp
        self._header.ssrc           = ssrc

    @property
    def sequenceNumber(self):
        return self._header.sequenceNumber
    @sequenceNumber.setter
    def sequenceNumber(self, no):
        self._header.sequenceNumber = no

    @property
    def timestamp(self):
        return self._header.timestamp
    @timestamp.setter
    def timestamp(self, t):
        self._header.timestamp = t

    @property
    def buffer(self):
        return self._buf

    def _setFixedValues(self):
        self._header.version   = 2
        self._header.padding   = 0
        self._header.extension = 0
        self._header.csrcCount = 0
        self._header.marker    = 0


class _PacketBuilder(object):
    def __init__(self, ssrc):
        self._seqNo     = 0
        self._timestamp = 0
        self._header    = _RtpHeader(_G711_PAYLOAD_TYPE, self._seqNo, self._timestamp, ssrc)

    def build(self, payload):
        self._header.sequenceNumber = self._seqNo
        self._header.timestamp      = self._timestamp
        packet = self._header.buffer + payload
        self._seqNo     += 1
        self._timestamp += len(payload) # given that one sample is represented as one byte (which is true for G.711)
        return packet


class _Reader(object):
    def __init__(self, pathToDataFile, chunkSize, dataQueue):
        assert chunkSize > 0
        self._pathToDataFile = pathToDataFile
        self._chunkSize = chunkSize
        self._dataQueue = dataQueue
        self._doRun = True

    def run(self):
        while self._doRun:
            lastChunk = bytes()
            lastChunk = self._readFile(lastChunk)

    def initShutdown(self):
        self._doRun = False

    def _readFile(self, lastChunk):
        with open(self._pathToDataFile, 'rb') as inputStream:

            wasIncomplete = len(lastChunk) < self._chunkSize
            if wasIncomplete:
                restCount = self._chunkSize - len(lastChunk)
                restChunk = inputStream.read(restCount)
                chunk = lastChunk + restChunk
                if len(chunk) < self._chunkSize:
                    return chunk
                self._dataQueue.put(chunk)

            while self._doRun:
                chunk = inputStream.read(self._chunkSize)
                if len(chunk) < self._chunkSize:
                    return chunk
                self._dataQueue.put(chunk)


class _Sender(object):
    def __init__(self, target):
        self._dest = (target.host, target.port)
        self._socket = socket.socket(type=socket.SOCK_DGRAM)

    def send(self, buf):
        self._socket.sendto(buf, self._dest)


class _QueueSender(object):
    def __init__(self, senders, dataQueue, packetBuilder, chunkSize):
        self._senders       = senders
        self._dataQueue     = dataQueue
        self._packetBuilder = packetBuilder
        self._doRun         = True
        self._nextTick      = None
        self._tick          = chunkSize/8000 # für G.711, 8kHz

    def run(self):
        while True:
            try:
                chunk = self._dataQueue.get(timeout=1)
            except Empty:
                if self._doRun:
                    continue
                else:
                    break
            self._send(chunk)
            self._dataQueue.task_done()

    def initShutdown(self):
        self._doRun = False

    def _send(self, chunk):
        packet = self._packetBuilder.build(chunk)

        if self._nextTick is None:
            self._nextTick = time.monotonic() + self._tick
            self._sendNowToAll(packet)
            return

        waitFor = self._nextTick - time.monotonic()
        assert waitFor >= 0, 'cannot send fast enough'
        time.sleep(waitFor)
        self._sendNowToAll(packet)
        self._nextTick += self._tick

    def _sendNowToAll(self, chunk):
        for sender in self._senders:
            sender.send(chunk)


_Target = namedtuple('_Target', ('host', 'port'))

def _createTarget(s, defaultPort):
    parts = s.split(':')
    assert len(parts) <= 2, 'invalid format for a target: {}'.format(s)
    if len(parts) == 1:
        return _Target(parts[0], defaultPort)
    else:
        return _Target(parts[0], int(parts[1]))

def _createTargets(targetDefs, defaultPort):
    return [_createTarget(t, defaultPort) for t in targetDefs]

def _parseCmdLine():
    parser = ArgumentParser(description=textwrap.dedent('''
        The program reads the given raw audio data file and sends RTP packets of the given size
        to the target hosts (can be more than one). When tha data file is exhausted, it starts
        reading from the front again until Ctrl-C is pressed.
        The data file must be raw G.711 data, 8kHz.
        Given a sound file (e.g. wav or ogg), a raw G.711 file at 8kHz can be generated using gstreamer
        in the following way:
        $ gst-launch-1.0 filesrc location=<path to sound file> ! decodebin ! audioconvert ! audio/x-raw, format=S16LE, layout=interleaved ! \\
            audioresample ! audio/x-raw, rate=8000 ! alawenc ! filesink location=g711.raw
        Alternatively you can also stream the data files stored under /data/AVS for some radio channel;
        they are stored in the required format.
    '''),
    formatter_class=RawDescriptionHelpFormatter)
    parser.add_argument('--target',           action='append', dest='targets', required=True, help='destination in the form host[:port]; can be given multiple times')
    parser.add_argument('--port', '-p',       type=int, metavar='N',           default=8001,  help='destination port if not specified in target definition; default: 8001')
    parser.add_argument('--file', '-f',                 metavar='PATH',        required=True, help='path to the data file')
    parser.add_argument('--ssrc',             type=int, metavar='N',           required=True, help='SSRC to be used')
    parser.add_argument('--chunk-size', '-s', type=int, metavar='N',           default=320,   help='size of the playload to use for each rtp packet; default: 320')
    options = parser.parse_args()
    assert len(options.targets) > 0, 'Need at least one target definition'
    options.targets=_createTargets(options.targets, options.port)
    del options.port
    return options


def _main():
    options = _parseCmdLine()
    dataQueue = Queue(maxsize=25)

    reader = _Reader(options.file, options.chunk_size, dataQueue)
    readerThread = Thread(target=reader.run)

    senders = [_Sender(target) for target in options.targets]
    queueSender = _QueueSender(senders, dataQueue, _PacketBuilder(options.ssrc), options.chunk_size)
    senderThread = Thread(target=queueSender.run)

    with ExitStack() as stack:
        stack.enter_context(startThreadContext(queueSender, senderThread))
        stack.enter_context(startThreadContext(reader, readerThread))
        while True:
            try:
                time.sleep(0.5)
            except KeyboardInterrupt:
                break


if __name__ == '__main__':
    _main()