"""
script to do a basic holistic test of the cowmesh network,
testing internet connection speed, and transfer speed,
at and between all nodes in the mesh
"""
import os
import time
import asyncio, asyncssh, sys
import subprocess
import json
import paramiko

PROJECT_PATH = os.path.abspath(os.path.dirname(__file__))
test_img_path = os.path.join(PROJECT_PATH, "test.png")

SECRETS_PATH = os.path.join(PROJECT_PATH, "secrets.json")
with open(SECRETS_PATH, 'r') as f:
    SECRETS = json.loads(f.read())

base_node = "janastunuc"  # nuc in jaaga

nodes = [
    "janastunuc",   # nuc in jaaga
    "solipi",       # pi in guard
    "writer",       # pi in new-gazebo
]


# url for downloading test.png from the internet
internet_url = "https://canalswans.commoninternet.net/test.png"

# boolean to prepare test, or just run it
PREPARE_TEST = False

if PREPARE_TEST:
    # copy image to internet
    print("COPYING TEST IMG TO INTERNET DOWNLOADABLE LOCATION")
    internet_scp_cmd = "scp -i ~/.ssh/do_rsa2 {} root@canalswans.commoninternet.net:/srv/canalswans.commoninternet.net/test.png".format(test_img_path)
    os.system(internet_scp_cmd)

    # first put the image on every node without testing the time (so that the test is prepared to run)
    for name, from_node_vals in SECRETS["HOST_INFO"].items():
        print("PREPARING NODE {}".format(name))
        ip = from_node_vals["ip"]
        dir_path = from_node_vals["node_path"]
        file_path = os.path.join(dir_path, "test.png")
        user = from_node_vals["user"]
        password = from_node_vals["password"]
        mkdir_cmd = "ssh {user}@{ip} 'mkdir -p {dir_path}'".format(user=user, ip=ip, dir_path=dir_path)
        os.system(mkdir_cmd)
        scp_cmd = "scp {test_img_path} {user}@{ip}:{file_path}".format(
            test_img_path=test_img_path,
            ip=ip,
            user=user,
            file_path=file_path
        )
        os.system(scp_cmd)

        myconn = paramiko.SSHClient()
        myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        myconn.connect(ip, username=user, password=password)
        remote_cmd = 'mkdir -p {dir_path}'.format(dir_path=dir_path)
        (stdin, stdout, stderr) = myconn.exec_command(remote_cmd)
        output = str(stdout.read())
        print("output: {}".format(output))
        print(("errors: {}".format(stderr.read())))

        scp_cmd = "scp {test_img_path} {user}@{ip}:{file_path}".format(
         test_img_path=test_img_path,
         ip=ip,
         user=user,
         file_path=file_path
        )
        (stdin, stdout, stderr) = myconn.exec_command(scp_cmd)
        output = str(stdout.read())
        print("output: {}".format(output))
        print(("errors: {}".format(stderr.read())))

        myconn.close()

results = {}

async def run_test():
    # run the test on each node
    node_dict = SECRETS["HOST_INFO"]
    for node in nodes:
        from_name = node
        from_node_vals = node_dict[node]
        from_ip = from_node_vals["ip"]
        from_dir_path = from_node_vals["node_path"]
        from_file_path = os.path.join(from_dir_path, "test.png")
        from_user = from_node_vals["user"]
        from_password = from_node_vals["password"]
        for to_name in nodes:
            to_node_vals = node_dict[to_name]
            if from_name == to_name:
                print("skip self")
                continue

            print("** running test.png transfer test from {} to {}".format(from_name, to_name))

            to_user = to_node_vals["user"]
            to_ip = to_node_vals["ip"]
            to_dir_path = to_node_vals["node_path"]
            to_file_path = os.path.join(to_dir_path, "download.png")
            to_password = to_node_vals["password"]

            scp_cmd = "scp -i {key_path} {from_file_path} {to_user}@{to_ip}:{to_file_path}".format(
                key_path="$HOME/.ssh/janastu",
                to_user=to_user,
                to_ip=to_ip,
                to_file_path=to_file_path,
                from_file_path=from_file_path
            )

            start = time.time()

            myconn = paramiko.SSHClient()
            myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            myconn.connect(from_ip, username=from_user, password=from_password)
            print(scp_cmd)
            (stdin, stdout, stderr) = myconn.exec_command(scp_cmd)
            exit_status = stdout.channel.recv_exit_status()
            output = str(stdout.read())
            error = str(stderr.read())
            print("output: {}".format(output))
            print(("errors: {}".format(stderr.read())))

            result_key = "{} -> {}".format(from_name, to_name)
            if exit_status == 0:
                print("success: {}".format(output))
                end = time.time()
                elapsed = (end - start)
                print("{}: {} seconds".format(result_key, elapsed))
                results[result_key] = elapsed
            else:
                print("error: {}".format(error))
                results[result_key] = "error: {}".format(error)
                return

            # scp_with_time_cmd = "set -e; /usr/bin/time -f '%e' {}".format(scp_cmd)
            # print("running: {}".format(scp_with_time_cmd))


def run_laptop_test():
    from_name = "laptop"
    # also measure transfers from laptop
    node_dict = SECRETS["HOST_INFO"]
    for to_name, to_node_vals in node_dict.items():
        print("** running test.png transfer test from laptop to {}".format(to_name))

        to_user = to_node_vals["user"]
        to_ip = to_node_vals["ip"]
        to_dir_path = to_node_vals["node_path"]
        to_file_path = os.path.join(to_dir_path, "download.png")

        scp_cmd = "scp -i {key_path} {from_file_path} {to_user}@{to_ip}:{to_file_path}".format(
            key_path="$HOME/.ssh/janastu",
            to_user=to_user,
            to_ip=to_ip,
            to_file_path=to_file_path,
            from_file_path=test_img_path
        )

        print(scp_cmd)

        start = time.time()
        result = subprocess.run(scp_cmd, shell=True)

        if result.returncode == 0:
            end = time.time()
            print("stdout: {}".format(result.stdout))
            elapsed = (end - start)
            result_key = "{}->{}".format(from_name, to_name)
            print("{}: {} seconds".format(result_key, elapsed))
            results[result_key] = elapsed
        else:
            print("stderr: {}".format(result.stderr))
            print('Program exited with status %d' % result.returncode,
                  file=sys.stderr)
            results[result_key] = "error: {}".format(result.stderr)
            return


try:
    # asyncio.get_event_loop().run_until_complete(run_test())
    run_laptop_test()

    print("** final results **")
    file_size = 13476
    for test_name, result in results.items():
        speed = file_size / float(result) / 1000.0
        print("{}: {:.2f} seconds, {:.2f} mbps".format(test_name, result, speed))

except (OSError, asyncssh.Error) as exc:
    sys.exit('SSH connection failed: ' + str(exc))