#!/usr/bin/env python3
#
# Convert comamnd-line options to a test specification

import optparse
import pscheduler
import sys
import datetime

logger = pscheduler.Log(prefix='test-throughput', quiet=True)

if len(sys.argv) > 1:

   # Args are on the command line
   args = sys.argv[1:]

else:

   # Args are in a JSON array on stdin
   json_args = pscheduler.json_load(exit_on_error=True)
   args = []

   if not isinstance(json_args,list):
      pscheduler.fail("Invalid JSON for this operation")
   for arg in json_args:
      if not ( isinstance(arg, str)
               or isinstance(arg, int)
               or isinstance(arg, float) ):
         pscheduler.fail("Invalid JSON for this operation")
   args = [ str(arg) for arg in json_args ]



# Gargle the arguments

opt_parser = pscheduler.FailingOptionParser(epilog=
"""Examples:

  task throughput --dest ps.example.com
      Measure througput from here to ps.example.com

  task throughput --source ps.example.org --dest ps.example.com
      Measure througput between ps.example.org and ps.example.com

  task throughput --bandwidth 384M --dest ps.example.com
      Limit bandwidth to 384 Mb per second

  task throughput --parallel 4 --dest ps.example.com
      Run four parallel streams
"""
                                            )

opt_parser.add_option("-s", "--source",
                      help="Sending host",
                      action="store", type="string",
                      dest="source")

opt_parser.add_option("--source-node",
                      help="pScheduler node on sending host, if different",
                      action="store", type="string",
                      dest="source_node")

opt_parser.add_option("-d", "--dest", "--destination",
                      help="Receiving host",
                      action="store", type="string",
                      dest="destination")

opt_parser.add_option("--dest-node",
                      help="pScheduler node on receiving host, if different",
                      action="store", type="string",
                      dest="dest_node")

opt_parser.add_option("-t", "--duration",
                      help="Total runtime of test",
                      action="store", type="string",
                      dest="duration")

opt_parser.add_option("-i", "--interval",
                      help="How often to report results (internally, results still reported in aggregate at end)",
                      action="store", type="string",
                      dest="interval")

opt_parser.add_option("--link-rtt",
                      help="Approximate link round-trip time (ISO8601 or integer ms)",
                      action="store", type="string",
                      dest="link_rtt")

opt_parser.add_option("-P", "--parallel",
                      help="How many parallel streams to run during the test",
                      action="store", type="int",
                      dest="parallel")

opt_parser.add_option("-u", "--udp",
                      help="Use UDP instead of TCP testing",
                      action="store_true", 
                      dest="udp")

opt_parser.add_option("-b", "--bandwidth",
                      help="Bandwidth to rate limit the test to, supports SI notation such as 1G",
                      action="store", type="string",
                      dest="bandwidth")

opt_parser.add_option("--bandwidth-strict",
                      help="Never go faster than --bandwidth, even to make up for lost time.",
                      action="store_true",
                      dest="bandwidth_strict")

opt_parser.add_option("--burst-size",
                      help="Limit bursts of packets to this number",
                      action="store", type="string",
                      dest="burst_size")

opt_parser.add_option("--fq-rate",
                      help="Fair queueing rate, supports SI notation such as 1G",
                      action="store", type="string",
                      dest="fq_rate")

opt_parser.add_option("-w", "--window-size",
                      help="TCP window (buffer) size to use for the test, supports SI notation such as 64M",
                      action="store", type="string",
                      dest="window_size")

opt_parser.add_option("-m", "--mss",
                      help="TCP maximum segment size",
                      action="store", type="int",
                      dest="mss")

opt_parser.add_option("-l", "--buffer-length",
                      help="length of the buffer to read/write from",
                      action="store", type="int",
                      dest="buffer_length")

opt_parser.add_option("--ip-tos",
                      help="IP type-of-service octet (integer)",
                      action="store", type="int",
                      dest="ip_tos")

opt_parser.add_option("--ip-version",
                      help="Specificy which IP version to use, 4 or 6",
                      action="store", type="int",
                      dest="ip_version")

opt_parser.add_option("-B", "--local-address",
                      help="Use this as a local address for control and tests",
                      action="store", type="string",
                      dest="local_address")

opt_parser.add_option("-O", "--omit",
                      help="Number of seconds to omit from the start of the test",
                      action="store", type="string",
                      dest="omit")

opt_parser.add_option("--no-delay",
                      help="Set TCP no-delay flag, disables Nagle's algorithm",
                      action="store_true",
                      dest="no_delay")

opt_parser.add_option("--congestion",
                      help="Set TCP congestion control algorithm",
                      action="store", type="string",
                      dest="congestion")

opt_parser.add_option("--zero-copy",
                      help="Use a 'zero copy' method of sending data",
                      action="store_true",
                      dest="zero_copy")

opt_parser.add_option("--flow-label",
                      help="Set the IPv6 flow label, implies --ip-version 6",
                      action="store", type="int",
                      dest="flow_label")

opt_parser.add_option("--client-cpu-affinity",
                      help="Set's the sending side's CPU affinity",
                      action="store", type="int",
                      dest="client_cpu_affinity")

opt_parser.add_option("--server-cpu-affinity",
                      help="Set's the receiving's side's CPU affinity",
                      action="store", type="int",
                      dest="server_cpu_affinity")

opt_parser.add_option("--single-ended",
                      help="Run a test directly to a host without pscheduler.",
                      action="store_true",
                      dest="single_ended")

opt_parser.add_option("--single-ended-port",
                      help="Run a test directly to a given port.",
                      action="store", type="int",
                      dest="single_ended_port")

opt_parser.add_option("--reverse",
                      help="Reverses the direction of the test.",
                      action="store_true",
                      dest="reverse")

opt_parser.add_option("--reverse-connections",
                      help="Make connections from destination to source where possible.",
                      action="store_true",
                      dest="reverse_connections")

opt_parser.add_option("--loopback",
                      help="Run both client and server side in a loopback test.",
                      action="store_true",
                      dest="loopback")

(options, remaining_args) = opt_parser.parse_args(args)

if len(remaining_args) != 0:
   pscheduler.fail("Unusable arguments: %s" % " ".join(remaining_args))


result = { }
schema = pscheduler.HighInteger(1)

if options.source is not None:
   result['source'] = options.source

if options.source_node is not None:
   result['source-node'] = options.source_node
   
if options.destination is not None:
   result['dest'] = options.destination

if options.dest_node is not None:
   result['dest-node'] = options.dest_node

if options.duration is not None:
   duration = options.duration

   # convert epoch seconds to is8601
   if duration.isdigit():
      delta = datetime.timedelta(seconds=int(duration))
      duration = pscheduler.timedelta_as_iso8601(delta)

   result['duration'] = duration

if options.interval is not None:

   interval = options.interval

   # convert epoch seconds to is8601
   if interval.isdigit():
      delta = datetime.timedelta(seconds=int(interval))
      interval = pscheduler.timedelta_as_iso8601(delta)

   result['interval'] = interval

if options.link_rtt is not None:
   rtt_value = options.link_rtt
   try:
      result['link-rtt'] = "PT%sS" % (int(rtt_value) / 1000.0)
   except ValueError:
      result['link-rtt'] = rtt_value
   schema.set(3)


if options.parallel is not None:
   result['parallel'] = options.parallel

if options.udp:
   result['udp'] = True

if options.bandwidth is not None:
   try:
      result['bandwidth'] = pscheduler.si_as_number(options.bandwidth)
   except ValueError as ex:
      pscheduler.fail("Invalid value \"%s\" for bandwidth: %s" % (options.bandwidth, ex))

if options.bandwidth_strict is not None:
   result['bandwidth-strict'] = options.bandwidth_strict
   schema.set(4)

if options.burst_size is not None:
   try:
      result["burst-size"] = pscheduler.si_as_number(options.burst_size)
      schema.set(4)
   except ValueError as ex:
      pscheduler.fail("Invalid value \"%s\" for burst size: %s" % (options.window_size, ex))

if options.fq_rate is not None:
   try:
      result['fq-rate'] = pscheduler.si_as_number(options.fq_rate)
      schema.set(7)
   except ValueError as ex:
      pscheduler.fail("Invalid value \"%s\" for fq-rate: %s" % (options.fq_rate, ex))

if options.window_size is not None:
   try:
      result["window-size"] = pscheduler.si_as_number(options.window_size)
   except ValueError as ex:
      pscheduler.fail("Invalid value \"%s\" for window size: %s" % (options.window_size, ex))

if options.mss is not None:
   result["mss"] = pscheduler.si_as_number(options.mss)

if options.buffer_length is not None:
   result["buffer-length"] = pscheduler.si_as_number(options.buffer_length)

if options.ip_version is not None:
   result["ip-version"] = options.ip_version

if options.local_address:
   result["local-address"] = options.local_address

if options.omit:
   omit = options.omit

   # convert epoch seconds to is8601
   if omit.isdigit():
      delta = datetime.timedelta(seconds=int(omit))
      omit  = pscheduler.timedelta_as_iso8601(delta)

   result["omit"] = omit
   

if options.ip_tos:
   result["ip-tos"] = options.ip_tos

if options.no_delay:
   result["no-delay"] = True


if options.congestion:

   # In schemas 1-3, this value was a hard-coded list.  In schema 4,
   # it is no longer validated and is passed directly to the tool.  We
   # don't have to bump the schema for the old, hard-coded values.
   if options.congestion not in [ "bbr", "bic", "cubic", "htcp",
                                  "reno", "vegas", "westwood", "yeah" ]:
      schema.set(4)

   result["congestion"] = options.congestion


if options.zero_copy:
   result["zero-copy"] = True

if options.flow_label:
   result["flow-label"] = int(options.flow_label)
   result["ip-version"] = 6 # flow label only works on ipv6

if options.client_cpu_affinity != None:
   result["client-cpu-affinity"] = options.client_cpu_affinity

if options.server_cpu_affinity != None:
   result["server-cpu-affinity"] = options.server_cpu_affinity

if options.single_ended:
   result["single-ended"] = True
   schema.set(2)

if options.single_ended_port:
   result["single-ended-port"] = options.single_ended_port
   schema.set(2)

if options.reverse:
   result["reverse"] = True

if options.reverse_connections:
   result["reverse-connections"] = True
   schema.set(6)

if options.loopback:
   result["loopback"] = True
   schema.set(5)
       
result["schema"] = schema.value()

logger.debug("cli-to-spec -> %s" % result)

pscheduler.succeed_json(result)
