#!<python_interpreter_path>

import paramiko
import json
import os
import paramiko
from scp import SCPClient
import argparse


LOCAL_PROJECT_PATH = "<project_path>"
config_path = f"{LOCAL_PROJECT_PATH}/scripts/remote_exe/remote_configs.json"
with open(config_path, "r") as f:
    remote_configs = json.load(f)


def remote_execute(cur_round, bench_path, run_flags, target):
    remote_config = remote_configs.get(target)
    remote_test_dir = remote_config["remote_test_dir"]

    local_bin_path = f"{bench_path}/bin/linked_{cur_round}.out"
    local_log_path = f"{bench_path}/perf_jsons/perf_stats_avg_linked_{cur_round}.json"
    local_perf_script = f"{LOCAL_PROJECT_PATH}/scripts/remote_exe/{remote_config["perf_script"]}"

    remote_log_path = local_log_path.replace(
        LOCAL_PROJECT_PATH, remote_test_dir)
    remote_bin_path = local_bin_path.replace(
        LOCAL_PROJECT_PATH, remote_test_dir)
    remote_perf_script = f"{remote_test_dir}/scripts/perf.sh"

    key_path = os.path.expanduser(remote_config["ssh_key"])

    # SCP send binary to target remote machine
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(remote_config["remote_host"], username=remote_config["remote_user"],
                       pkey=paramiko.RSAKey.from_private_key_file(key_path))

        cmd = f"""
mkdir -p {"/".join(remote_bin_path.split("/")[:-1])}
mkdir -p {"/".join(remote_log_path.split("/")[:-1])}
mkdir -p {"/".join(remote_perf_script.split("/")[:-1])}
"""
        print(cmd)

        stdin, stdout, stderr = client.exec_command(cmd)
        exit_code = stdout.channel.recv_exit_status()

        with SCPClient(client.get_transport()) as scp:
            scp.put(local_bin_path, recursive=True,
                    remote_path=remote_bin_path)

        with SCPClient(client.get_transport()) as scp:
            scp.put(local_perf_script, recursive=True,
                    remote_path=remote_perf_script)

        client.close()
        print(f"Binary sent to remote machine: {target}")

    except Exception as e:
        print(f"ERROR: {e}")
        return

    # Run remote binary
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(remote_config["remote_host"], username=remote_config["remote_user"],
                       pkey=paramiko.RSAKey.from_private_key_file(key_path))
        run_flags = run_flags.replace(LOCAL_PROJECT_PATH, remote_test_dir)

        cmd = f"""
bash {remote_perf_script} {remote_bin_path} {run_flags}
"""

        stdin, stdout, stderr = client.exec_command(cmd)
        exit_code = stdout.channel.recv_exit_status()

        client.close()
        print(f"Executed remote binary {remote_bin_path} test on: {target}")

    except Exception as e:
        print(f"ERROR: {e}")
        return

    # Pull remote performance json back
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(remote_config["remote_host"], username=remote_config["remote_user"],
                       pkey=paramiko.RSAKey.from_private_key_file(key_path))

        with SCPClient(client.get_transport()) as scp:
            scp.get(remote_path=remote_log_path, local_path=local_log_path)

        client.close()
        print(f"Performance json pulled down from: {target}")

    except Exception as e:
        print(f"ERROR: {e}")
        return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Main script to run AutoPass")

    parser.add_argument("--curround", type=str,
                        required=True, help="Currend round number")
    parser.add_argument("--benchpath", type=str,
                        required=True, help="benchmark's path")
    parser.add_argument(
        "--runflags", type=str, required=True, help="Args of program"
    )
    parser.add_argument(
        "--target", type=str, required=True, help="Build target."
    )

    args = parser.parse_args()
    cur_round = args.curround
    bench_path = args.benchpath
    runflags = args.runflags
    target = args.target

    # Simple test
    # cur_round = 1
    # bench_path = f"{LOCAL_PROJECT_PATH}/benchmarks/coremark"
    # runflags = " needoutput "
    # target = "pi"

    remote_execute(cur_round=cur_round, bench_path=bench_path,
                   run_flags=runflags, target=target)

    exit(0)
