#!/usr/bin/env python3
import argparse
import os
import sys


def main():
    args = handle_args()
    cmd = build_cmd(args)
    # Run the command
    print(f"Running command: {cmd}")
    result = os.system(cmd)
    if result != 0:
        print("")
        print("Removing problematic SSH key and trying again..")
        remove_ssh_key_cmd = (
            'ssh-keygen -f "${HOME}/.ssh/known_hosts" -R "[localhost]:1535"'
        )
        os.system(remove_ssh_key_cmd)
        print(f'Running command "{cmd}"')
        result = os.system(cmd)


def handle_args():
    help_string = (
        "This script copies files to the Q7S as long as port forwarding is active.\n"
    )
    help_string += (
        "You can set up port forwarding with "
        '"ssh -L 1535:192.168.133.10:22 <eive-flatsat-ip>" -t /bin/bash'
    )
    parser = argparse.ArgumentParser(description=help_string)
    # Optional arguments
    parser.add_argument(
        "-r", "--recursive", dest="recursive", default=False, action="store_true"
    )
    parser.add_argument(
        "-t",
        "--target",
        help="Target destination. If files are copied to Q7S, will be /tmp by default. "
        "If files are copied back to host, will be current directory by default",
        default="",
    )
    parser.add_argument("-P", "--port", help="Target port", default=1535)
    parser.add_argument(
        "-i",
        "--invert",
        default=False,
        action="store_true",
        help="Copy from Q7S to host instead. Always copies to current directory.",
    )
    # Positional argument(s)
    parser.add_argument(
        "source", help="Source files to copy or target files to copy back to host"
    )
    return parser.parse_args()


def build_cmd(args):
    # Build run command
    cmd = "scp "
    if args.recursive:
        cmd += "-r "
    target = args.target
    if args.invert and target == "":
        target = "."
    elif target == "":
        target = f"/tmp"
    if args.invert:
        cmd += f"-P {args.port} root@localhost:{args.source} {target}"
    else:
        cmd += f"-P {args.port} {args.source} root@localhost:{target}"
    if args.target:
        cmd += args.target
    return cmd


if __name__ == "__main__":
    main()