#!/usr/bin/python
#
# Run an nuttcp test
#

import datetime
import logging
import json
import pscheduler
import re
import shutil
import sys
import time
import nuttcp_parser
import traceback
import ipaddr
import nuttcp_utils
from nuttcp_defaults import *



# track when this run starts
start_time = datetime.datetime.now()

logger = pscheduler.Log(prefix='tool-nuttcp', quiet=True)

logger.debug("starting nuttcp tool")

# parse JSON input
input = pscheduler.json_load(exit_on_error=True)

logger.debug("Input is %s" % input)

try:
    participant = input['participant']
    participant_data = input['participant-data']
    test_spec = input['test']['spec']
    duration = pscheduler.iso8601_as_timedelta(input['schedule']['duration'])
    limits_passed = input.get('limits-passed', [])
except KeyError as e:
    pscheduler.fail("Missing required key in run input: %s" % e)
except:
    pscheduler.fail("Error parsing run input: %s" % sys.exc_info()[0])

if len(participant_data) != 2:
    pscheduler.fail("nuttcp requires exactly 2 participants, got %s" % (len(participant_data)))


config = nuttcp_utils.get_config()
nuttcp_cmd  = config["nuttcp_cmd"]

server_port = participant_data[1]['server_port']
data_port_start = participant_data[1]['data_port_start']

def run_client():    

    logger.debug("Waiting %s sec for server on other side to start" % DEFAULT_WAIT_SLEEP)

    time.sleep(DEFAULT_WAIT_SLEEP) #wait for server to start on other side

    nuttcp_args = [ nuttcp_cmd ]

    nuttcp_args.append('-p')
    nuttcp_args.append(data_port_start)

    nuttcp_args.append('-P')
    nuttcp_args.append(server_port)


    # Determine if we need to specify a specific ip-version if we have enough info to do so intelligently
    ip_version = test_spec.get('ip-version', None)
    source = test_spec.get('source', None)

    if ip_version is None and source is not None:
        source_ip, dest_ip = pscheduler.ip_normalize_version(source, test_spec['dest'])         
        if source_ip is not None and dest_ip is not None:
            ip_version = pscheduler.ip_addr_version(source_ip)[0]
            
    if ip_version:
        nuttcp_args.append('-%s' % ip_version)


    # duration
    test_duration = test_spec.get('duration')
    if test_duration:
        delta = pscheduler.iso8601_as_timedelta(test_duration)
        test_duration = int(pscheduler.timedelta_as_seconds(delta))
    else:
        test_duration = DEFAULT_DURATION

    nuttcp_args.append('-T')
    nuttcp_args.append(test_duration)

    # big list of optional arguments to nuttcp, map from test spec
    if test_spec.has_key('interval') and test_spec['interval'] != None:
        nuttcp_args.append('-i')
        delta = pscheduler.iso8601_as_timedelta(test_spec['interval'])
        nuttcp_args.append(int(pscheduler.timedelta_as_seconds(delta)))

    if test_spec.has_key('parallel') and test_spec['parallel'] != None:
        nuttcp_args.append('-N')
        nuttcp_args.append(test_spec['parallel'])

    if test_spec.has_key('window-size') and test_spec['window-size'] != None:
        nuttcp_args.append('-w')
        nuttcp_args.append(test_spec['window-size'])

    if test_spec.has_key('mss') and test_spec['mss'] != None:
        nuttcp_args.append('-M')
        nuttcp_args.append(test_spec['mss'])

    if test_spec.has_key('bandwidth') and test_spec['bandwidth'] != None:
        nuttcp_args.append('-R')
        # nuttcp specifies bandwidth in Kbps, so need to convert here
        nuttcp_args.append(test_spec['bandwidth'] / 1000)
    else:

        # Find the largest bandwidth in each participant's passed
        # limits and take the minimum of those values.  This will be
        # the maximum bandwidth allowed by any participant.

        logger.debug("Finding maximum bandwidth")

        max_bandwidths = []

        for passed_set in limits_passed:
            set_max = -1
            logger.debug("Limit set:")
            for passed in passed_set:
                try:
                    new_max = pscheduler.si_as_number(passed['bandwidth']['range']['upper'])
                    logger.debug("  Limit says %d", new_max)
                    set_max = max(set_max, new_max)
                except KeyError:
                    logger.debug("  No upper bandwidth limit specified")
                    continue

            if set_max > -1:
                max_bandwidths.append(set_max)

        if max_bandwidths:
            max_bandwidth = sorted(max_bandwidths)[0]
            logger.debug("Limiting bandwidth to %d", max_bandwidth)
            nuttcp_args.append('-R')
            # nuttcp specifies bandwidth in Kbps, so need to convert here
            nuttcp_args.append(max_bandwidth / 1000)


    if test_spec.has_key('udp'):
        nuttcp_args.append('-u')
        nuttcp_args.append('-j') # add jitter reporting

    if test_spec.has_key('buffer-length') and test_spec['buffer-length'] != None:
        nuttcp_args.append('-l')
        nuttcp_args.append(test_spec['buffer-length'])

    if test_spec.has_key('ip-tos') and test_spec['ip-tos'] != None:
        nuttcp_args.append('-c')
        # The 'T' suffix treats the TOS value as the whoe octet
        # instead of just the DSCP.
        nuttcp_args.append("%sT" % (str(test_spec['ip-tos'])))

    if test_spec.get('reverse'):
        nuttcp_args.append('-F')

    # who to connect to
    destination = test_spec['dest']

    affinity = test_spec.get('client-cpu-affinity')
    if affinity == None:
        affinity = pscheduler.source_affinity(destination, ip_version=ip_version)

    if affinity != None:
        nuttcp_args.insert(0, 'numactl')
        nuttcp_args.insert(1, '-N')
        nuttcp_args.insert(2, affinity)


    nuttcp_args.append(destination)

    # join and run_program want these all to be string types, so
    # just to be safe cast everything in the list to a string
    nuttcp_args = map(lambda x: str(x), nuttcp_args)
      
    logger.debug("Running command: %s" % " ".join(nuttcp_args))

    try:
        status, stdout, stderr = pscheduler.run_program(nuttcp_args)
    except Exception as e:
        logger.error("nuttcp failed to complete execution: %s" % e)
        pscheduler.succeed_json({"succeeded": False,
                                 "error": "The nuttcp command failed during execution. See server logs for more details."})


    return _make_result(" ".join(nuttcp_args), status, stdout, stderr)


def run_server():

    #init command
    nuttcp_args = [ nuttcp_cmd, '-S', '-1', '--nofork']

    nuttcp_args.append("-p")
    nuttcp_args.append(data_port_start)

    nuttcp_args.append("-P")
    nuttcp_args.append(server_port)

    # Determine if we need to specify a specific ip-version if we have enough info to do so intelligently
    ip_version = test_spec.get('ip-version', None)
    source = test_spec.get('source', None)
    if ip_version is None and source is not None:
        source_ip, dest_ip = pscheduler.ip_normalize_version(source, test_spec['dest'])         
        if source_ip is not None and dest_ip is not None:
            ip_version = pscheduler.ip_addr_version(source_ip)[0]
            
    if ip_version:
        nuttcp_args.append('-%s' % ip_version)

    # try to grab our default affinity if one wasn't passed in
    affinity = test_spec.get('server-cpu-affinity')
    if affinity == None:
        # look up what interface we're going to be receiving on        
        interface = pscheduler.address_interface(test_spec['dest'], ip_version=ip_version)
        if interface:
            affinity = pscheduler.interface_affinity(interface)

    if affinity != None:
        nuttcp_args.insert(0, 'numactl')
        nuttcp_args.insert(1, '-N')
        nuttcp_args.insert(2, affinity)

    nuttcp_args = map(lambda x: str(x), nuttcp_args)
    logger.debug("Running command: %s" % " ".join(nuttcp_args))

    try:
        status, stdout, stderr = pscheduler.run_program(nuttcp_args,
                                                        timeout = pscheduler.timedelta_as_seconds(duration - (datetime.datetime.now() - start_time)),
                                                        timeout_ok = True)
    except Exception as e:
        logger.error("nuttcp failed to complete execution: %s" % e);
        pscheduler.succeed_json({"succeeded": False,
                                 "error": "The nuttcp command failed during execution. See server logs for more details."})


    # in nuttcp, the client always reports the results even in reverse mode so 
    # we never need to parse the server's output
    return _make_result(" ".join(nuttcp_args), status, stdout, stderr, parse=False)



def _make_result(command_line, status, stdout, stderr, parse=True):
    logger.debug("Stdout = %s" % stdout)
    logger.debug("Stderr = %s" % stderr)
    
    if status:
        pscheduler.succeed_json({"succeeded": False,
                                 "error": "nuttcp returned an error: %s\n%s" % (stdout, stderr)})

    lines = stdout.split("\n")    
    logger.debug("Lines are %s " % lines)

    results = {}
    if parse:
        results = nuttcp_parser.parse_output(lines)

    results['diags'] = "%s\n\n%s\n%s" % (command_line, stdout, stderr)

    return results


#determine whether we are the client or server mode for nuttcp
results = {}
try:
    if participant == 0:
        results = run_client()
    elif participant == 1:
        results = run_server()
    else:
        pscheduler.fail("Invalid participant.")
except Exception as ex:
    _, _, ex_traceback = sys.exc_info()
    if ex_traceback is None:
        ex_traceback = ex.__traceback__
    tb_lines = [ line.rstrip('\n') for line in
                 traceback.format_exception(ex.__class__, ex, ex_traceback)]
    logger.debug(tb_lines)
    logger.error("Exception %s" % ex);

logger.debug("Results: %s" % results)

pscheduler.succeed_json(results)
