from random import choice
from string import ascii_lowercase

from ...Logging.Logger import log, Severity
from ...Machine.PaedMLMachineCollection import MACHINE_COLLECTION
from ...SetupSteps.Playbooks.Playbook import FileCleanupTask, PlaybookCredentialCollection
from ...SetupSteps.SetupStep import SetupStep


class PerformOrderedServerReboots(SetupStep):
    def run(self, file_cleanup_stack : list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        log('Staging opsi server reboot...', Severity.INFO)
        temp_shellscript = f'/tmp/rs_1_{"".join(choice(ascii_lowercase) for i in range(5))}'
        temp_helpscript = f'/tmp/rs_2_{"".join(choice(ascii_lowercase) for i in range(5))}'
        MACHINE_COLLECTION.backup.execute_shell_command('touch', temp_shellscript)
        MACHINE_COLLECTION.backup.execute_shell_command('touch', temp_helpscript)
        MACHINE_COLLECTION.backup.execute_shell_command('chmod', '700', temp_shellscript)
        MACHINE_COLLECTION.backup.execute_shell_command('chmod', '700', temp_helpscript)


        MACHINE_COLLECTION.backup.write_file(temp_shellscript, f'''#! /bin/bash
        while nc -z {MACHINE_COLLECTION.server.get_hostname()} 445; do
            sleep 1
        done
        while ! nc -z {MACHINE_COLLECTION.server.get_hostname()} 445; do
            sleep 1
        done
        reboot
        '''.encode('utf-8'))

        MACHINE_COLLECTION.backup.write_file(temp_helpscript, f'''#! /bin/bash
        nohup {temp_shellscript} &
        '''.encode('utf-8'))

        # Perform evil wizardry to work around our own safety measures
        stdout, stderr, retval = MACHINE_COLLECTION.backup.execute_shell_command(temp_helpscript, non_blocking=True)
        if retval != 0:
            log(f'Failed to run reboot script: {stderr}', Severity.ERROR)
            return False

        log('Going down for reboot NOW')
        MACHINE_COLLECTION.server.execute_shell_command('reboot')
        return True