from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from pssh.clients.miko.single import SSHClient as MikoClient
from pssh.utils import load_private_key

import os
import sys
from datetime import datetime

from pssh.clients.native.single import SSHClient


user = sys.argv[1]
host = sys.argv[2]
pkey_path = os.path.expanduser("~/.ssh/id_rsa")

cmd = "echo 'hello world'"


paramiko_times = []


miko_client = MikoClient(host, user=user, pkey=load_private_key(pkey_path))


times = []
for i in range(9):
    start = datetime.now()
    channel, _, stdout, stderr, _ = miko_client.exec_command(cmd)
    stdout = list(stdout)
    channel.close()
    total = datetime.now() - start
    paramiko_times.append(total)
    assert stdout[0].strip() == "hello world"


psshrun_times = []

client = SSHClient(host, user=user, pkey=pkey_path)

for i in range(7):
    start = datetime.now()
    res = client.run_command(cmd)
    channel = res[0]
    stdout = list(res[2])
    client.close_channel(channel)
    _run = datetime.now() - start
    assert stdout[0].strip() == "hello world"
    psshrun_times.append(_run)


psshopen_times = []
psshcmd_times = []
psshtotal_times = []
psshclose_times = []

for i in range(7):
    start = datetime.now()
    channel = client.open_session()
    chan_open = datetime.now() - start
    start = datetime.now()
    client.execute(cmd, channel=channel)
    execute = datetime.now() - start
    start = datetime.now()
    stdout = "\n".join(client.read_output(channel))
    _read = datetime.now() - start
    start = datetime.now()
    client.close_channel(channel)
    psshclose = datetime.now() - start
    
    assert stdout.strip() == "hello world"

    psshopen_times.append(chan_open)
    psshcmd_times.append(execute)
    psshtotal_times.append(_read)
    psshclose_times.append(psshclose)

def fmt_times(name, _times):
    # Convert to ms
    times = [t.total_seconds()*1000 for t in _times]
    return "{:<20}  {:8.2f}    {:8.2f}    {:8.2f}".format(
        name, min(times), max(times), sum(times) / len(times)
    )

print("                      min [ms]    max [ms]    avg [ms]")
print(fmt_times("paramiko", paramiko_times))
print(fmt_times("pssh run_command", psshrun_times))
print(fmt_times("pssh execute", psshtotal_times))
print(fmt_times("     open", psshopen_times))
print(fmt_times("     cmd", psshcmd_times))
print(fmt_times("     close", psshclose_times))