#!/usr/bin/python3
#
# Copyright (c) 2014-2017 - Adjacent Link LLC, Bridgewater, New Jersey
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of Adjacent Link LLC nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#

from argparse import ArgumentParser
import otestpoint.interface.discovery_pb2 as discovery_pb2
import sqlite3
import zmq
import time
import socket
import struct
import fcntl
import sys

def get_ip_address(ifname):
    # http://code.activestate.com/recipes/439094/
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    return socket.inet_ntoa(fcntl.ioctl(
        s.fileno(),
        0x8915,  # SIOCGIFADDR
        struct.pack('256s', ifname[:15].encode() if sys.version_info >= (3,0) else ifname[:15])
    )[20:24])

argumentParser = ArgumentParser()


argumentParser.add_argument('-i',
                            '--interface',
                            dest='interface',
                            type=str,
                            default='lo',
                            help='server interface address. [default: %(default)s]')

argumentParser.add_argument('--port-discover',
                            dest='port-discover',
                            type=int,
                            default=9001,
                            help='server discover port. [default: %(default)s]')

argumentParser.add_argument('--port-publish',
                            dest='port-publish',
                            type=int,
                            default=9002,
                            help='server publish port. [default: %(default)s]')


argumentParser.add_argument('-r',
                            '--rate',
                            dest='rate',
                            type=int,
                            default=1,
                            help='playback rate in timestamp per second. [default: %(default)s]')

argumentParser.add_argument('recorder-file',
                            nargs='+',
                            type=str,
                            help="otestpoint-recorder file.")

ns = argumentParser.parse_args()

args = vars(ns)

entries = []

timestamps = set()
probes = set()

class Entry:
    pass

for filename in args['recorder-file']:
    entry = Entry()
    entry.name = filename
    entry.fd = open(filename,'rb')
    entry.conn = sqlite3.connect(filename+'.db')

    # build a set of all timestamps
    sql='SELECT time,probe FROM probes ORDER BY ROWID ASC'

    for row in entry.conn.execute(sql):
        timestamps.add(int(row[0]))
        probes.add(row[1])

    entry.cursor = entry.conn.cursor()

    entries.append(entry)


address = get_ip_address(args['interface'])

context = zmq.Context()
publish = context.socket(zmq.PUB)
publish.bind('tcp://%s:%d' % (address,args['port-publish']))

server = context.socket(zmq.REP)
server.bind('tcp://%s:%d' % (address,args['port-discover']))

poller = zmq.Poller()
poller.register(server, zmq.POLLIN)

timestamps_sorted = sorted(timestamps)

last_timestamp_publish = time.time()

timeout = 1/float(args['rate'])

while True:
    try:
        socks = dict(poller.poll(timeout * 1000))
    except:
        break

    if socks.get(server) == zmq.POLLIN:
        request = discovery_pb2.DiscoveryRequest()

        request.ParseFromString(server.recv())

        response = discovery_pb2.DiscoveryResponse()

        if request.type == discovery_pb2.DiscoveryRequest.TYPE_DISCOVERY:

            response.type = discovery_pb2.DiscoveryResponse.TYPE_DISCOVERY

            response.discovery.publish = 'tcp://%s:%d' % (address,args['port-publish'])

            for probe in sorted(probes):
                response.discovery.names.append(probe)
        else:
            response.type = discovery_pb2.DiscoveryResponse.TYPE_ERROR

            response.error.what = "Unkown discovery type"

        server.send(response.SerializeToString())

    current_time = time.time()

    if current_time - timeout >= last_timestamp_publish:
        last_timestamp_publish = current_time

        timestamp = timestamps_sorted.pop(0)

        for entry in entries:

            sql='SELECT offset,size,probe FROM probes WHERE time == ?;'

            entry.cursor.execute(sql,[timestamp])

            while True:
                row = entry.cursor.fetchone()

                if row != None:
                    offset,size,probe = row

                    entry.fd.seek(offset)

                    msg=entry.fd.read(size)
                    publish.send_multipart([probe.encode(),msg])
                else:
                    break

        # all done
        if not timestamps_sorted:
            break

for entry in entries:
    entry.fd.close()
    entry.cursor.close()
    entry.conn.close()
