import os
import socket
from time import sleep
from typing import Tuple

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

class RebootFirewall(SetupStep):
    # Constants
    REBOOT_WAIT_MAX_SECONDS = 480
    PORT = 80
    SHORT_TIMEOUT = 5  # For detecting reboot
    LONG_TIMEOUT = 10  # For detecting recovery

    def __init__(self):
        super().__init__()

    def _is_port_available(self, timeout: int) -> bool:
        """Check if the port is available (i.e., firewall is down or up)."""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)

        try:
            sock.connect((MACHINE_COLLECTION.firewall.get_hostname(), self.PORT))
            sock.close()
            return True
        except TimeoutError:
            return False
        except (OSError, ConnectionRefusedError):
            sleep(timeout)
            return False

    def _wait_for_condition(self, condition_func, timeout: int, max_seconds: int) -> bool:
        """Wait for a condition to be met (e.g., firewall down or up)."""
        tries_left = max_seconds / timeout
        while tries_left > 0:
            if condition_func(timeout):
                return True
            tries_left -= 1
            log(f"Still waiting...", Severity.TRACE)
        return False

    def run(self, file_cleanup_stack : list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        log("Rebooting firewall...")
        stdout, stderr, retval = MACHINE_COLLECTION.firewall.execute_shell_command('reboot')

        if len(stderr) != 0:
            log(stderr, Severity.ERROR)
            return False

        # Wait for firewall to go down (port unavailable)
        if not self._wait_for_condition(
            lambda t: not self._is_port_available(t),
            self.SHORT_TIMEOUT,
            self.REBOOT_WAIT_MAX_SECONDS
        ):
            log("Firewall did not reboot.", Severity.ERROR)
            return False

        log("Waiting for Firewall to come back up.", Severity.DEBUG)

        # Wait for firewall to come back up (port available)
        if not self._wait_for_condition(
            lambda t: self._is_port_available(t),
            self.LONG_TIMEOUT,
            self.REBOOT_WAIT_MAX_SECONDS
        ):
            log("Firewall did not come back online.", Severity.ERROR)
            return False

        log("Firewall rebooted and is online.", Severity.DEBUG)
        return True
