#!/usr/bin/env python3
#
# Run a test.
#

import datetime
import json
import pscheduler
import random
import socket
import sys
import time
import traceback

log = pscheduler.Log(prefix="simplestreamer", quiet=True)

input = pscheduler.json_load(exit_on_error=True);

log.debug("Input is %s", input)

#
# Validate the input
#

if not 'participant-data' in input:
    pscheduler.fail("No participant data.")





# Perform the test

# Ideally, the timeout would be for the whole operation (establishing
# a connection and sending/receiving the data), but that
# implementation is more complex than necessary for this example.
# This makes a very rough approximation 

start_time = datetime.datetime.now()
try:
    timeout = pscheduler.iso8601_as_timedelta(input['test']['spec']['timeout'])
except KeyError:
    timeout = pscheduler.iso8601_as_timedelta('PT5S')
end_time = datetime.datetime.now() + timeout

def time_left():
    return pscheduler.timedelta_as_seconds(end_time - datetime.datetime.now())


random.seed()

# Force a hard failure if one is indicated
try:
    fail_prob = float(input['test']['spec']['fail'])
    fail_rand = random.random()
    if fail_rand < fail_prob:
        results = {
            'succeeded': False,
            'diags': ("Random %f on prob. %f" % (fail_rand, fail_prob)),
            'error': 'Randomly-induced failure',
            'result': None
            }
        pscheduler.succeed_json(results)
except KeyError:
    pass  # Skip it if not in the input


# Both participants need these.

ip_version = input['test']['spec'].get('ip-version')

source_host = input['test']['spec'].get('source', pscheduler.api_local_host())
dest_host = input['test']['spec']['dest']

log.debug(f'''Normalizing {source_host} -> {dest_host}''')
send_ip, recv_ip = pscheduler.ip_normalize_version(source_host, dest_host, ip_version=ip_version)
log.debug(f'''Got {send_ip} -> {recv_ip}''')
if send_ip is None or recv_ip is None:
    pscheduler.succeed_json({
        "succeeded": False,
        "error": f'''Unable to find common IP version between {source_host} and {dest_host}.'''
    })



participant = input['participant']


try:
    start_at = input['schedule']['start']
    log.debug("Sleeping until %s", start_at)
    pscheduler.sleep_until(start_at)
    log.debug("Starting")
except KeyError:
    pscheduler.fail("Unable to find start time in input")


if participant == 0:

    #
    # Sender
    #

    # Dawdle if directed.

    try:
        dawdle = pscheduler.timedelta_as_seconds(
            pscheduler.iso8601_as_timedelta(input['test']['spec']['dawdle']) )
        dawdle *= random.random()
        log.debug("Dawdling %s", dawdle)
        time.sleep(dawdle)
    except KeyError:
        dawdle = 0

    # Connect and send the material

    succeeded = False
    error = None

    try:
        material = input['test']['spec']['test-material']
    except KeyError:
        material = "Data sent at " + str(datetime.datetime.now())


    source_address = (send_ip, 0)
    server_address = (recv_ip, input['participant-data'][1]['listen-port'])

    try_num = 1
    succeeded = False
    diags = []

    sock = None

    while datetime.datetime.now() < end_time:

        time_remaining = time_left()
        log.debug("Timeout is %s", time_remaining)

        log.debug("Connecting %s -> %s", str(source_address), str(server_address))
        try:
            sock = socket.create_connection(server_address, time_remaining, source_address)
            if time_left() <= 0:
                log.debug("Timed out")
                diags.append("Ran out of time trying to connect.")
                break
            log.debug("Connected")
            sock.sendall(material.encode())
            log.debug("Sent")
            succeeded = True
            error = None
            diags.append("Try %d succeeded." % (try_num))
            break
        except Exception as ex:
            error = ex.args[1] if isinstance(ex, socket.gaierror) else str(ex)
            error = "Failed to connect: %s" % (error)
            diags.append("Try %d failed: %s" % (try_num, error))
            log.debug(error)
            sleep_time = time_left() * 0.10
            if sleep_time > 0:
                log.debug("Sleeping %s", sleep_time)
                time.sleep(sleep_time)
        finally:
            if sock is not None:
                sock.close()

        try_num += 1

    results = {
        'succeeded': succeeded,
        'diags': "\n".join(diags),
        'error': error,
        'result': {
            'dawdled': pscheduler.timedelta_as_iso8601(
                pscheduler.seconds_as_timedelta(dawdle)),
            'sent': material
            } if succeeded else None
        }

elif participant == 1:

    #
    # Dest
    #

    diags = []
    succeeded = False
    full = ''
    error = None

    server_address = (recv_ip, input['participant-data'][1]['listen-port'])
    diags.append("Listening on %s" % str(server_address))

    try:
        socket_family = pscheduler.ip_addr_version(recv_ip, family=True)[0]
        diags.append("Socket family is %s" % str(socket_family))
        if socket_family is None:
            raise socket.gaierror("Unable to determine address family.")

        sock = socket.socket(socket_family, socket.SOCK_STREAM)
        sock.settimeout(time_left())
        sock.bind(server_address)
        sock.listen(1)
        connection, client_address = sock.accept()
        connection.settimeout(time_left())

        while True:
            data = connection.recv(1024)
            if data:
                full += data.decode("ascii")
            else:
                break

        connection.close()
        succeeded = True

    except socket.timeout:
        error = "Client never connected."

    except socket.gaierror as ex:
        (err_no, err_str) = ex.args
        error = "Unable to set up listener: %s" % (err_str)

    except Exception as ex:
        error = "Receiver exception: " + str(ex) \
            + " " + traceback.format_exc()

    results = {
        'succeeded': succeeded,
        'diags': "\n".join(diags),
        'error': error if not succeeded else None,
        'result': {
            'received': full,
            'elapsed-time':  pscheduler.timedelta_as_iso8601(
                datetime.datetime.now() - start_time)
            } if succeeded else None
        }

else:

    pscheduler.fail("Invalid participant.")


pscheduler.succeed_json(results)
