#!/usr/bin/env python3

import argparse, subprocess, os, sys, logging, time
from enum import Enum


def exit_on_failure(func):
    def wrapper(*args, **kwargs):
        return_code = func(*args, **kwargs)
        if return_code:
            logging.error("Error: " + func.__name__ + " " + str(args) + " " +
                          str(kwargs) + " exited with return code " +
                          str(return_code))
            cont = input("Continue (y/N)? ")
            if cont != 'y':
                exit(1)

    return wrapper


class NaoVersion(Enum):
    V5 = 0
    V6 = 1


def base_dir():
    return os.path.dirname(os.path.abspath(__file__))


def determine_nao_version(nao_number: int):
    if nao_number >= 20:
        return NaoVersion.V6
    else:
        return NaoVersion.V5


def get_nao_address(nao_number: int, ethernet: bool = True):
    if ethernet:
        return "10.1.YOUR_TEAM_NUMBER_HERE." + str(nao_number + 10)
    else:
        return "10.0.YOUR_TEAM_NUMBER_HERE." + str(nao_number + 10)


@exit_on_failure
def nao_copy_ssh_id(address):
    subprocess.run(["cp", base_dir() + "/ssh_key", base_dir() + "/files"])
    ssh_copy_id_command = [
        "sshpass", "-p", "nao", "ssh-copy-id", "-o",
        "UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "-o",
        "LogLevel=quiet", "-o", "ConnectTimeout=5", "-i",
        base_dir() + "/files/ssh_key", "nao@" + address
    ]
    returncode = subprocess.run(ssh_copy_id_command).returncode
    subprocess.run(["rm", "-f", base_dir() + "/files/ssh_key"])
    return returncode


@exit_on_failure
def nao_cmd(address, command):
    ssh_command = [
        "ssh", "-o", "UserKnownHostsFile=/dev/null", "-o",
        "StrictHostKeyChecking=no", "-o", "LogLevel=quiet", "-o",
        "ConnectTimeout=5", "-l", "nao", "-i",
        base_dir() + "/ssh_key", address, command
    ]
    return subprocess.run(ssh_command).returncode


@exit_on_failure
def nao_cp(sources, destination):
    if type(sources) != list:
        sources = [sources]

    scp_command = [
        "scp", "-o", "UserKnownHostsFile=/dev/null", "-o",
        "StrictHostKeyChecking=no", "-o", "LogLevel=quiet", "-i",
        base_dir() + "/ssh_key", "-r", *sources, destination
    ]
    return subprocess.run(scp_command).returncode


def nao_reboot(address):
    reboot_command = [
        "ssh", "-o", "UserKnownHostsFile=/dev/null", "-o",
        "StrictHostKeyChecking=no", "-o", "LogLevel=quiet", "-o",
        "ConnectTimeout=5", "-l", "nao", "-i",
        base_dir() + "/ssh_key", address, "systemctl reboot"
    ]
    return subprocess.run(reboot_command).returncode


if __name__ == "__main__":
    param_parser = argparse.ArgumentParser(
        "Gammary Script to setup a NAO for hulks code base")
    param_parser.add_argument(
        "nao_number", type=int, help="the number the NAO should get (considered v6 if nao_number >= 20)")
    param_parser.add_argument(
        "--address", "-a", help="connect to the NAO via a special ip address")
    param_parser.add_argument(
        "--wait",
        "-w",
        action="store_true",
        default=False,
        help= "Not available for V6! Wait until the NAO is reachable after reboot instead of quitting"
    )
    param_parser.add_argument(
        "--sysroot-only",
        "-s",
        action="store_true",
        default=False,
        help="only install the sysroot")
    param_parser.add_argument(
        "--skip-sysroot",
        action="store_true",
        default=False,
        help="skip installing the sysroot")
    param_parser.add_argument(
        "--purge", "-p", help="remove all existing sysroots on the NAO")

    args = param_parser.parse_args()

    logging.basicConfig(
        format='%(levelname)s:%(message)s', level=logging.DEBUG)
    logging.info("Running Gammaray Script")

    nao_version = determine_nao_version(args.nao_number)

    if nao_version == NaoVersion.V5:
        logging.info("Determined Nao version V5")
        logging.info("Running V5 Gammaray")
        old_gammaray_cmd = [base_dir() + "/lib/gammarayV5"]
        old_gammaray_cmd.extend(sys.argv[1:])
        subprocess.run(old_gammaray_cmd)
        exit(0)

    if args.address:
        nao_address = args.address
    else:
        nao_address = get_nao_address(args.nao_number)

    if not args.sysroot_only:
        logging.info("Copy SSH-Key")
        nao_copy_ssh_id(nao_address)

    if not args.skip_sysroot:
        logging.info("Copy sysroot")
        nao_cp(base_dir() + "/../toolchain/sysroot.tar.bz2",
               "nao@" + nao_address + ":")
        logging.info("Extracting sysroot")
        nao_cmd(nao_address, "tar xf sysroot.tar.bz2")
        nao_cmd(nao_address, "rm sysroot.tar.bz2")
    else:
        logging.info("Skipping sysroot...")

    if args.sysroot_only:
        exit(0)

    logging.info("Set Hostname")
    nao_cmd(nao_address,
            "hostnamectl set-hostname YOUR_ROBOT_NAME_HERE" + str(args.nao_number))

    logging.info("Copy files_hulks")
    nao_cp(base_dir() + "/files_hulks/bin", "nao@" + nao_address + ":.local")
    nao_cp(base_dir() + "/files_hulks/etc", "nao@" + nao_address + ":.local")
    nao_cp(base_dir() + "/files_hulks/asound.state", "nao@" + nao_address + ":.config")
    nao_cp(base_dir() + "/files_hulks/.asoundrc", "nao@" + nao_address + ":~")

    logging.info("Set executable bits")
    nao_cmd(nao_address, "chmod u+x .local/bin/connman-hulks")
    nao_cmd(nao_address, "chmod u+x .local/bin/setNetwork")
    nao_cmd(nao_address, "chmod u+x .local/bin/alsa-restore")

    logging.info("Copy bashrc")
    nao_cp(base_dir() + "/files_hulks/bashrc_v6", "nao@" + nao_address + ":.profile")

    logging.info("Set nao ip")
    nao_cmd(
        nao_address, "sed -e s/{wifi_ip}/" + get_nao_address(
            args.nao_number, False) + "/ -e s/{eth_ip}/" + get_nao_address(
                args.nao_number) + "/ -i /home/nao/.local/etc/connman-hulks.yaml")

    logging.info("Copy services")
    nao_cmd(nao_address, "mkdir -p .config/systemd/user")
    nao_cp([
        base_dir() + "/files_hulks/connman-hulks.service",
        base_dir() + "/files_hulks/lola-hulks.service",
        base_dir() + "/files_hulks/hulk.service",
        base_dir() + "/files_hulks/alsa-restore.service"
    ], "nao@" + nao_address + ":.config/systemd/user/")

    logging.info("Activate services")
    nao_cmd(nao_address, "systemctl --user enable lola-hulks.service")
    nao_cmd(nao_address, "systemctl --user enable connman-hulks.service")
    nao_cmd(nao_address, "systemctl --user enable hulk.service")
    nao_cmd(nao_address, "systemctl --user enable alsa-restore.service")

    logging.info("Add robocup.conf")
    nao_cmd(nao_address, "touch robocup.conf")

    logging.info("Reboot NAO")
    nao_reboot(nao_address)
