"""
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


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


node_dict = {
    "jagga-nuc": {
        "ip": "10.56.10.247",
        "user": "cow",
        "node_path": "/datadrive/data/cowmesh-network-test"
    },
    "redcottage-pi": {
        "ip": "10.56.148.122",
        "user": "admin",
        "node_path": "/home/admin/cowmesh-network-test"
    },
    "newgazebo-pi": {
        "ip": "10.56.11.111",
        "user": "pi",
        "node_path": "/home/pi/cowmesh-network-test"
    },
}

# 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 node_dict.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"]
        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)

results = {}

async def run_test():
    # run the test on each node
    for from_name, from_node_vals in node_dict.items():
        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"]
        for to_name, to_node_vals in node_dict.items():
            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")

            async with asyncssh.connect(from_ip, username=from_user) as conn:

                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
                )

                print(scp_cmd)

                start = time.time()
                result = await conn.run(scp_cmd)

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