import subprocess
import time

import yaml
from tqdm import tqdm


def main():
    print("Retrieving hosts")
    with open("ansible/inventory.yaml", "r") as f:
        all_hosts = yaml.load(f, yaml.SafeLoader)
    fedora_hosts = all_hosts["fedora"]["hosts"].keys()

    with open("update.log", "w") as log_file:
        for _, host in enumerate(tqdm(fedora_hosts, desc="Running system updates")):
            log_file.write(f"Updating {host}\n")
            log_file.flush()
            subprocess.run(
                ["ssh", host, "dnf", "upgrade", "-y"],
                stdout=log_file,
                stderr=log_file,
                check=True,
            )
            log_file.write(f"Rebooting {host}\n")
            log_file.flush()
            subprocess.run(
                ["ssh", host, "reboot"],
                stdout=log_file,
                stderr=log_file,
                check=True,
            )
            time.sleep(5)  # wait for reboot to take effect
            booted = False
            max_attempts = 5  # seconds
            cur_attempts = 0  # seconds
            while cur_attempts > max_attempts or not booted:
                try:
                    subprocess.run(
                        ["ssh", host, "echo"],
                        stdout=log_file,
                        stderr=log_file,
                        check=True,
                        timeout=2,
                    )
                    booted = True
                except Exception as e:
                    cur_attempts += 1
                    log_file.write(f"Waiting for {host} to reboot\n")
                    log_file.flush()
                    time.sleep(5)

            if cur_attempts >= max_attempts:
                exit(1)


if __name__ == "__main__":
    main()