import socket

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

_SSH_DIR = '/root/.ssh'
_KNOWN_HOSTS = '/root/.ssh/known_hosts'
_AUTHORIZED_KEYS = '/root/.ssh/authorized_keys'
_ID_KEY = '/root/.ssh/id_ed25519'
_ID_KEY_PUB = '/root/.ssh/id_ed25519.pub'


class SetupSshTrust(SetupStep):
    """
    Ensures every machine (server, backup, firewall) can SSH to every other
    machine without any interactive prompt:
      - known_hosts on each machine contains fingerprints for all aliases
        (FQDN, short hostname, IP) of all machines.
      - authorized_keys on each machine contains the public keys of all machines.
    """

    def _resolve_aliases(self, machines: list[MachineInterface]) -> list[list[str]] | None:
        """Returns a list of alias lists (FQDN, short name, IP) for each machine,
        or None if any hostname cannot be resolved."""
        machine_aliases: list[list[str]] = []
        for machine in machines:
            fqdn = machine.get_hostname()
            short = machine.get_hostname_short()
            try:
                ip = socket.gethostbyname(fqdn)
            except socket.gaierror as e:
                log(f"Could not resolve IP for {fqdn}: {e}", Severity.ERROR)
                return None
            seen: set[str] = set()
            aliases: list[str] = []
            for alias in [fqdn, short, ip]:
                if alias not in seen:
                    seen.add(alias)
                    aliases.append(alias)
            machine_aliases.append(aliases)
            log(f"Aliases for {short}: {aliases}", Severity.DEBUG)
        return machine_aliases

    def _collect_host_keys(self, machine: MachineInterface, aliases: list[str]) -> list[str] | None:
        """Reads all sshd host public keys from a machine and returns known_hosts lines.
        Runs directly on the machine to avoid any local host-key verification."""
        stdout, stderr, retval = machine.execute_shell_command(
            'sh', '-c', 'cat /etc/ssh/ssh_host_*_key.pub'
        )
        if retval != 0 or not stdout.strip():
            log(f"Could not read sshd host keys on {aliases[0]}: {stderr.strip()}", Severity.ERROR)
            return None
        aliases_str = ','.join(aliases)
        lines = []
        for line in stdout.strip().splitlines():
            parts = line.strip().split()
            if len(parts) >= 2:
                lines.append(f"{aliases_str} {parts[0]} {parts[1]}")
        if not lines:
            log(f"No host keys found on {aliases[0]}", Severity.ERROR)
            return None
        log(f"Collected {len(lines)} host key(s) from {aliases[0]}.", Severity.DEBUG)
        return lines

    def _build_known_hosts(
        self,
        machines: list[MachineInterface],
        machine_aliases: list[list[str]],
    ) -> bytes | None:
        """Collects host keys from all machines and returns combined known_hosts content."""
        all_lines: list[str] = []
        for machine, aliases in zip(machines, machine_aliases):
            lines = self._collect_host_keys(machine, aliases)
            if lines is None:
                return None
            all_lines.extend(lines)
        return ('\n'.join(all_lines) + '\n').encode('utf-8')

    def _deploy(self, machines: list[MachineInterface], known_hosts: bytes) -> bool:
        """Ensures each machine has an SSH key pair, then distributes known_hosts
        and authorized_keys to all machines. Returns False on any failure."""
        public_keys: list[str] = []
        for machine in machines:
            machine.mkdir(_SSH_DIR)
            if machine.check_if_file_exists(_ID_KEY):
                machine.delete(_ID_KEY)
            if machine.check_if_file_exists(_ID_KEY_PUB):
                machine.delete(_ID_KEY_PUB)
            log(f"Generating SSH key pair on {machine.get_hostname_short()}...", Severity.INFO)
            stdout, stderr, retval = machine.execute_shell_command(
                'ssh-keygen', '-t', 'ed25519', '-N', '', '-f', _ID_KEY, '-q'
            )
            if retval != 0:
                log(f"Failed to generate SSH key on {machine.get_hostname_short()}: {stderr}",
                    Severity.ERROR)
                return False
            log("Key generated.", Severity.DEBUG)
            pub_key = machine.read_file(_ID_KEY_PUB).decode('utf-8').strip()
            public_keys.append(pub_key)

        authorized_keys_content = ('\n'.join(public_keys) + '\n').encode('utf-8')

        for machine in machines:
            machine.mkdir(_SSH_DIR)
            machine.write_file(_KNOWN_HOSTS, known_hosts)
            machine.write_file(_AUTHORIZED_KEYS, authorized_keys_content)
            machine.execute_shell_command('chmod', '700', _SSH_DIR)
            machine.execute_shell_command('chmod', '600', _KNOWN_HOSTS)
            machine.execute_shell_command('chmod', '600', _AUTHORIZED_KEYS)
            log(f"SSH trust files written to {machine.get_hostname_short()}.", Severity.DEBUG)

        return True

    def run(self, file_cleanup_stack: list[FileCleanupTask], credentials: PlaybookCredentialCollection) -> bool:
        machines = [
            MACHINE_COLLECTION.server,
            MACHINE_COLLECTION.backup,
            MACHINE_COLLECTION.firewall,
        ]

        machine_aliases = self._resolve_aliases(machines)
        if machine_aliases is None:
            return False

        known_hosts = self._build_known_hosts(machines, machine_aliases)
        if known_hosts is None:
            return False

        if not self._deploy(machines, known_hosts):
            return False

        log("SSH trust established between all machines for all aliases.", Severity.INFO)
        return True
