#!/usr/bin/python3
import argparse
import sys
from colorama import Fore, Back, Style
import platform
import datetime
import subprocess
import getpass
import paramiko
import logging
import json
import secrets

VERSION = "0.1.3"

password: str = ""  # To be set by user
system_architecture: str = ""  # To be set by fact gathering
distribution: str = ""  # To be set by user

remote_folder = "/tmp/opensuse-anywhere/"
repo_link = ""  # To be set by user-ish

parser = argparse.ArgumentParser(
    prog="opensuse-anywhere",
    description="Install the openSUSE distributions on any running Linux system!",
    epilog=f"Developed by Alexander Johansen (alejoh.eu). This project is inspired by nixos-anywhere. You are running version {VERSION}."
)

parser.add_argument("host", help="what host to SSH into")
parser.add_argument("run", help="runs opensuse-anywhere on a host")

parser.add_argument(
    "--key", help="path to private SSH key to authenticate with.")
args = parser.parse_args(args=None if sys.argv[1:] else ['--help'])

# Print version info etc on run
print(f"{Fore.GREEN}opensuse-anywhere{Style.RESET_ALL} by Alexander Johansen (alejoh.eu), version {Fore.LIGHTBLUE_EX}{VERSION}{Style.RESET_ALL}, running on {platform.platform()} at {datetime.datetime.now(datetime.timezone.utc)} UTC")


def ask_for_password() -> str:
    '''
    All this function does is ask the user for the password for the root user on their host.

    Args:
        args: args from argparse

    Returns:
        str: password in string
    '''
    return getpass.getpass(
        f"{Back.GREEN}⌨️ 🔑 Please enter the password for the root user on {args.host}:{Style.RESET_ALL}")


def generate_password() -> str:
    '''
    Generates a password using the secrets module.

    Args:
        None

    Returns:
        str: generated password
    '''
    return secrets.token_urlsafe(16)

def determine_repo(arch: str, distribution: str) -> str:
    '''
    This function determines the repository link to use based on the distribution picked by the user and the remote host's architecture.

    Args:
        arch (str): the remote host's architecture
        distribution: the distribution the user picked
    Returns:
        str: link to use for repository
    '''
    if distribution == "leap":
        if arch == "aarch64":
            repo_link = f"https://download.opensuse.org/ports/{arch}/distribution/leap/15.6/repo/oss/"
        else:
            repo_link = f"https://download.opensuse.org/distribution/leap/15.6/repo/oss/"
    else:
        if arch == "aarch64":
            repo_link = f"https://download.opensuse.org/ports/{arch}/{distribution}/repo/oss/"
        else:
            repo_link = f"https://download.opensuse.org/{distribution}/repo/oss/"
    return repo_link


def determine_download_link(arch: str, repo_link: str) -> str:
    '''
    This function determines the root of the link used by curl to download the Linux kernel & initrd based on architecture and pre-existing repo link.

    Args:
        arch (str): The architecture of the remote system
        repo_link (str): The repo link to use as a base
    Returns:
        str: the download link
    '''

    if arch == "aarch64":
        download_link = repo_link + f"boot/{arch}/"
    else:
        download_link = repo_link + f"boot/{arch}/loader/"
    return download_link


def run_command(cmd: str) -> dict[str, str | int]:
    '''
    This function just runs a command via SSH, and returns a dict with the output (as a single line) + the return code of the command.

    Args:
        cmd (str): The command to run.

    Returns:
        dict[str, str | int]: A dictionary containing the output (output) and return code (return_code) of the executed command.
    '''
    # if args.key == "" or args.key == None:
    #    command = f"ssh root@{args.host} {cmd}"
    # else:
    #    command = f"ssh root@{args.host} -i {args.key} {cmd}"

    # cmdd = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    # stdout, stderr = cmdd.communicate()
    # return_code = cmdd.wait()
    # return return_code

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    if args.key == "" or args.key == None:
        client.connect(args.host, username="root", password=password)
    else:
        key = paramiko.Ed25519Key.from_private_key_file(args.key)
        client.connect(args.host, username="root", pkey=key)
    _stdin, _stdout, _stderr = client.exec_command(cmd, get_pty=True)
    output = str(_stdout.readline())
    return_code = _stdout.channel.recv_exit_status()
    client.close()
    return {"output": output, "return_code": return_code}


def check_host() -> bool:
    '''
    This function checks if the specified host is ready for booting the YaST installer.

    For example, it checks the following crucial stuff:
        Root access
        kexec support
        Curl

    Args:
        None

    Returns:
        bool: Whether the system is ready or not.
    '''
    # Define globals
    global system_architecture
    # Check remote system architecture
    arch_command = run_command(
        '[ "$(uname -m)" = "x86_64" ] && exit 0 || [ "$(uname -m)" = "aarch64" ] && exit 1 || [ "$(uname -m)" = "armv7l" ] && exit 2')
    if arch_command.get("return_code") == 0:
        print(
            Fore.BLUE + f"ℹ️  {args.host} is running on the x86_64 architecture." + Style.RESET_ALL)
        system_architecture = "x86_64"
    elif arch_command.get("return_code") == 1:
        print(Fore.YELLOW + f"⚠️ {args.host} is running on the aarch64 architecture. If this system is in the cloud, for example Hetzner, then this should be fine. However, if this is a Raspberry Pi, this program is unsupported, and you should just flash the SD card." + Style.RESET_ALL)
        system_architecture = "aarch64"
    elif arch_command.get("return_code") == 2:
        print(Fore.RED + f"⛔ {args.host} is running on the armv7l architecture, which has openSUSE images available, for the Raspberry Pi. The Raspberry Pi is not supported by opensuse-anywhere. You should just flash the SD card." + Style.RESET_ALL)
        return False
    # Check if we have root access
    if run_command('if [ $UID -eq 0 ]; then echo "Root!"; exit 0; else echo "Not root!"; exit 1; fi').get("return_code") == 0:
        print(Fore.GREEN +
              f"✅ We have root access on {args.host}." + Style.RESET_ALL)
    else:
        print(
            Fore.RED + f"⛔ We do not have root access on {args.host}." + Style.RESET_ALL)
        return False
    # Check if kexec is available
    if run_command('which kexec').get("return_code") == 0:
        print(Fore.GREEN +
              f"✅ kexec-tools is installed on {args.host}." + Style.RESET_ALL)
    else:
        print(Fore.RED + f"⛔ {args.host} does not have the binary {Style.RESET_ALL}{Fore.BLUE}kexec{Style.RESET_ALL}{Fore.RED} installed. This is usually provided via the {Style.RESET_ALL}{Fore.BLUE}kexec-tools{Style.RESET_ALL}{Fore.RED} package.")
        return False
    # Check if systemctl is available
    if run_command('which systemctl').get("return_code") == 0:
        print(Fore.GREEN +
              f"✅ systemctl is available on {args.host}." + Style.RESET_ALL)
    else:
        print(Fore.RED + f"⛔ systemctl is not available on {args.host}. Are you using another init system than systemd, like OpenRC? These are not yet supported. If there is any interest for OpenRC support, please create an issue on the repo and I will consider it." + Style.RESET_ALL)
        return False
    # Check if we have Curl installed
    if run_command('which curl').get("return_code") == 0:
        print(Fore.GREEN +
              f"✅ curl is installed on {args.host}." + Style.RESET_ALL)
    else:
        print(Fore.RED + f"⛔ {args.host} does not have {Style.RESET_ALL}{Fore.BLUE}curl{Style.RESET_ALL}{Fore.RED} installed. This is used for testing internet connectivity and downloading the MicroOS ISO image. Please install it.{Style.RESET_ALL}")
        return False
    # Check for Internet access
    if run_command("curl https://download.opensuse.org").get("return_code") == 0:
        print(
            Fore.GREEN + f"✅ {args.host} can reach openSUSE's Download service." + Style.RESET_ALL)
    else:
        print(Fore.RED + f"⛔ {args.host} is unable to reach openSUSE's Download service (download.opensuse.org). Please take a look into this." + Style.RESET_ALL)
        return False
    return True


def select_distro():
    '''
    This function asks the user which openSUSE distribution they would like to install.

    You have the options of:
    1. Tumbleweed
    2. Leap
    3. Slowroll

    Other distributions are not yet supported.

    Args:
        None
    Returns:
        None
    '''
    global distribution
    distribution = input(
        f"🖥️⌨️ Please enter the openSUSE distribution you would like to install (supported are tumbleweed, leap and slowroll):")
    if distribution == "slowroll":
        if system_architecture == "aarch64":
            print(f"{Fore.YELLOW}⚠️🖥️ Slowroll is not yet supported on {system_architecture}. You may want to select Tumbleweed.{Style.RESET_ALL}")
            sys.exit()
        print(f"{Fore.YELLOW}⚠️🖥️ Slowroll needs some extra steps after install. Please check the README for more information.{Style.RESET_ALL}")


def procedure() -> bool:
    '''
    This function is the one that does all the magic.
    '''
    # Define globals
    global ip_address
    global yast_password
    # Tell users on system that opensuse-anywhere is running
    if run_command("echo opensuse-anywhere is running on this system! It will reboot in a few minutes, so save your work. | wall").get("return_code") == 0:
        print(f"{Fore.GREEN}🔔✅ Notified active terminals on the system to save their work.{Style.RESET_ALL}")
    else:
        print(
            f"{Fore.RED}🔔⛔ Could not notify active terminals with wall.{Style.RESET_ALL}")
    # Create tmpdir folder
    if run_command(f"rm -rf {remote_folder} && mkdir {remote_folder}").get("return_code") == 0:
        print(f"{Fore.GREEN}📁✅ Created {remote_folder}.{Style.RESET_ALL}")
    else:
        print(f"{Fore.RED}⛔📁 Could not create {remote_folder}.{Style.RESET_ALL}")
        return False
    # Download openSUSE distribution on the remote system
    print(f"{Fore.BLUE}⬇️ℹ️ Downloading openSUSE {distribution.capitalize()} {system_architecture}. This may take a while, but rest assured, the script is still working.{Style.RESET_ALL}")
    repo_link = determine_repo(system_architecture, distribution)
    download_link = determine_download_link(
        system_architecture, repo_link)
    if run_command(f"curl {download_link}linux -L -o {remote_folder}linux && curl {download_link}initrd -L -o {remote_folder}initrd").get("return_code") == 0:
        print(f"{Fore.GREEN}⬇️✅ Downloaded openSUSE {distribution.capitalize()} {system_architecture} to {remote_folder}/extracted.{Style.RESET_ALL}")
    else:
        print(f"{Fore.RED}⬇️⛔ Could not download {distribution.capitalize()} {system_architecture} to {remote_folder}/extracted.{Style.RESET_ALL}")
        return False
    # Get networking info (chunky json)
    # Get default gateway
    default_gateway = run_command('ip route show default | awk \'{print $3}\'')
    real_default_gateway: str = ""
    if default_gateway.get("return_code") == 0:
        real_default_gateway = str(default_gateway.get("output"))
        print(f"{Fore.GREEN}🌐✅ Fetched {args.host}'s gateway.{Style.RESET_ALL}")
    else:
        print(f"{Fore.RED}🌐⛔ Could not fetch {args.host}'s gateway.{Style.RESET_ALL}")
        return False
    ip_address = run_command(
        'ip address show dev "$(ip route show default | awk \'{{print $5}}\')" | grep \'inet \' | awk \'{print $2}\'')
    if ip_address.get("return_code") == 0:
        print(f"{Fore.GREEN}🌐✅ Fetched {args.host}'s IP address.{Style.RESET_ALL}")
    else:
        print(
            f"{Fore.RED}🌐⛔ Could not fetch {args.host}'s IP address.{Style.RESET_ALL}")
        return False
    # Create a separate password for YaST for security reasons (this will be shown to the user)
    yast_password = generate_password()
    # Set up kexec
    kexec = run_command(f"kexec -l {remote_folder}linux --command-line=\"ifcfg=*={str(ip_address.get('output')).strip()},{real_default_gateway.strip()},9.9.9.11 149.112.112.11 ssh=1 ssh.password={yast_password} textmode=1 install={repo_link} self_update=0\" --initrd={remote_folder}initrd")
    if kexec.get("return_code") == 0:
        print(f"{Fore.GREEN}🖥️✅ Loaded kernel & initrd via kexec.{Style.RESET_ALL}")
    else:
        print(f"{Fore.RED}🖥️⛔ Something went wrong while running kexec: {str(kexec.get('output'))}{Style.RESET_ALL}")
        return False
    # Run kexec
    if run_command(f"systemctl kexec").get("return_code") == 0:
        print(
            f"{Fore.GREEN}🖥️✅ Rebooting into new kernel via systemd.{Style.RESET_ALL}")
    else:
        print(
            f"{Fore.RED}🖥️⛔ Could not reboot into new kernel via systemd.{Style.RESET_ALL}")
        return False
    return True


# "Main loop"
if args.run:
    print(Fore.CYAN + "Info: " + Style.RESET_ALL +
          "Checking if remote system is compatible.")
    if args.key == "" or args.key == None:
        password = ask_for_password()
    else:
        password = generate_password()
    if check_host() == False:
        print(Fore.RED + "❌ The remote host does not meet a required criteria. Please fix this, and try again." + Style.RESET_ALL)
        sys.exit(1)
    print(Fore.CYAN + "Info: " + Style.RESET_ALL +
          f"{args.host} passed the compatibility check. Proceeding.")
    select_distro()
    if procedure() == False:
        print(f"{Fore.RED}❌ Something went wrong with running the procedure on {args.host}.{Style.RESET_ALL}")
        sys.exit(1)
    else:
        print(f"{Fore.GREEN}🖥️✅ The openSUSE installer should be booting on {args.host} right now. SSH is enabled, so wait approx. 3 minutes and enter the following command in your terminal: {Style.RESET_ALL}{Fore.BLUE}ssh -X root@{args.host}{Style.RESET_ALL}{Fore.GREEN}. The password is {Style.RESET_ALL}{Fore.BLUE}{yast_password}{Style.RESET_ALL}{Fore.GREEN}. Enjoy!{Style.RESET_ALL}")
        sys.exit(0)
