#! /usr/bin/python3
# Copyright (c) 2014,2016,2019 - 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.
#
# See toplevel COPYING for more information.
#

from __future__ import absolute_import, division, print_function
import zmq
import uuid
import struct
import sys
import time
import six
from optparse import OptionParser
import otestpoint.interface.probereport_pb2
from otestpoint.interface import make_measurement_operator

usage = """%prog [OPTION]... ENDPOINT [PROBENAME]...

  ENDPOINT     OpenTestPoint publisher
               <hostname>:port | <IPv4>:port | [<IPv6>]:port

  PROBENAME    OpenTestPoint probe name. Not specifying any
               probe names means subscribe to all."""

description="""Subscribe to one, some or all probes published by ENDPOINT and
print OpenTestPoint probe information or binary probes in probe stream
format."""

epilog="""Probe stream format uses length prefix framing, where the length of
the message is output as an unsigned 32-bit integer value (4 bytes) in
network byte order preceding the output of the serialized message. All
probe messages are serialized OpenTestPopint::ProbeReport protobuf
messages. OpenTestPoint probe name format consists of a tree like
naming convention where each element in the name tree is separated by
a '.'. The more elements in a probe name, the more specific the probe
name. For example, a probe named A.B.C.D may belong to a family of
probes consisting of A.B.C.D, A.B.C.E and A.B.C.F. Specifying a probe
named A.B.C.D will only match that single probe. Specifying A.B.C will
match A.B.C.D, A.B.C.E and A.B.C.F. Likewise, specifying just A will
match all probes that start with A."""

optionParser = OptionParser(usage=usage,
                            description=description,
                            epilog=epilog)

optionParser.add_option("-b",
                        "--binary",
                        action="store_true",
                        dest="binary",
                        help="write binary messages to stdout")

optionParser.add_option("-t",
                        "--tally",
                        action="store_true",
                        dest="tally",
                        help="count probes and display output")

(options, args) = optionParser.parse_args()

if len(args) < 1:
  print("missing endpoint")
  exit(1)


context = zmq.Context()

subscriber = context.socket(zmq.SUB)

subscriber.setsockopt(zmq.IPV4ONLY,0)

subscriber.connect("tcp://%s" % args[0])

if len(args) > 1:
  for probe in args[1:]:
    if six.PY2:
      subscriber.setsockopt(zmq.SUBSCRIBE,
                            probe)
    else:
      subscriber.setsockopt_string(zmq.SUBSCRIBE,
                                   probe)

else:
  subscriber.setsockopt(zmq.SUBSCRIBE, b'')

tally = {}
nodes = set()

poller = zmq.Poller()

poller.register(subscriber, zmq.POLLIN)

nextTally=time.time() + 2

maxProbeLength = 0
maxNodeLength  = 5
operators = {}

try:
  while True:

    if options.tally:
      timeout = (nextTally - time.time()) * 1000
    else:
      timeout=None

    socks = dict(poller.poll(timeout))

    # Handle worker activity on backend
    if (subscriber in socks and socks[subscriber] == zmq.POLLIN):

      msgs = subscriber.recv_multipart()

      if options.binary:
        if six.PY2:
          sys.stdout.write(struct.pack("!L",len(msgs[1])));
          sys.stdout.write(msgs[1])
        else:
          sys.stdout.buffer.write(struct.pack("!L",len(msgs[1])));
          sys.stdout.buffer.write(msgs[1])
        sys.stdout.flush()
      else:
        report = otestpoint.interface.probereport_pb2.ProbeReport()
        report.ParseFromString(msgs[1])

        if options.tally:
          index =  msgs[0].rfind(b'.')
          probe = msgs[0][:index]
          node = msgs[0][index+1:]

          if probe not in tally:
            maxProbeLength = max(maxProbeLength,len(probe))
            tally[probe] = {}

          if node not in tally[probe]:
            maxNodeLength = max(maxNodeLength,len(node))
            tally[probe][node] = 1
          else:
            tally[probe][node] += 1

          nodes.add(node)

        else:
          print()
          print("[%s]" % report.timestamp, "%s/%d" % (report.tag,report.index),uuid.UUID(bytes=report.uuid))
          print(report.data.module, report.data.name,"v%d" % report.data.version,len(report.data.blob),"bytes")
          print(msgs[0].decode("utf-8"))

          if report.data.name not in operators:
            MeasurementOperator = make_measurement_operator(report.data.module,
                                                            report.data.name)

            if MeasurementOperator != None:
              operator = MeasurementOperator()
            else:
              operator = None

            operators[report.data.name] = operator
          else:
            operator = operators[report.data.name]

          if operator != None:
            measurement = operator.create(report.data.blob)
            print(operator.pprint(measurement))

    if options.tally and time.time() >= nextTally:
      print('\x1b[2J\x1b[H')
      probeFmt ="%%-%ds" % maxProbeLength
      nodeFmt="%%-%ds" % maxNodeLength
      print(probeFmt % "", end=' ')

      for node in sorted(nodes):
        print(nodeFmt % node.decode("utf-8"), end=' ')
      print()

      for probe in sorted(tally.keys()):
        print(probeFmt % probe.decode("utf-8"), end=' ')

        for node in sorted(nodes):
          if node in  tally[probe]:
            total = tally[probe][node]
          else:
            total = 0
          print(nodeFmt % total, end=' ')

        print()

      nextTally=time.time() + 2


except KeyboardInterrupt:
  pass

print(file=sys.stderr)
