from __future__ import annotations

import shlex
from random import uniform
from shutil import which
from socket import error as socket_error, gaierror
from time import sleep
from typing import IO, TYPE_CHECKING, Any, Iterable, Optional, Protocol, Tuple

import click
from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
from paramiko.agent import Agent
from typing_extensions import TypedDict, Unpack, override

from pyinfra import logger
from pyinfra.api.command import QuoteString, StringCommand
from pyinfra.api.exceptions import ConnectError
from pyinfra.api.util import get_file_io, memoize

from .base import BaseConnector, DataMeta
from .scp import SCPClient
from .ssh_util import get_private_key, raise_connect_error
from .sshuserclient import SSHClient
from .util import (
    CommandOutput,
    execute_command_with_sudo_retry,
    make_unix_command_for_host,
    read_output_buffers,
    run_local_process,
    write_stdin,
)

if TYPE_CHECKING:
    from pyinfra.api.arguments import ConnectorArguments


class ConnectorData(TypedDict):
    ssh_hostname: str
    ssh_port: int
    ssh_user: str
    ssh_password: str
    ssh_key: str
    ssh_key_password: str

    ssh_allow_agent: bool
    ssh_look_for_keys: bool
    ssh_forward_agent: bool

    ssh_config_file: str
    ssh_known_hosts_file: str
    ssh_strict_host_key_checking: str

    ssh_paramiko_connect_kwargs: dict

    ssh_connect_retries: int
    ssh_connect_retry_min_delay: float
    ssh_connect_retry_max_delay: float
    ssh_file_transfer_protocol: str


connector_data_meta: dict[str, DataMeta] = {
    "ssh_hostname": DataMeta("SSH hostname"),
    "ssh_port": DataMeta("SSH port"),
    "ssh_user": DataMeta("SSH user"),
    "ssh_password": DataMeta("SSH password"),
    "ssh_key": DataMeta("SSH key filename"),
    "ssh_key_password": DataMeta("SSH key password"),
    "ssh_allow_agent": DataMeta(
        "Whether to use any active SSH agent",
        True,
    ),
    "ssh_look_for_keys": DataMeta(
        "Whether to look for private keys",
        True,
    ),
    "ssh_forward_agent": DataMeta(
        "Whether to enable SSH forward agent",
        False,
    ),
    "ssh_config_file": DataMeta("SSH config filename"),
    "ssh_known_hosts_file": DataMeta("SSH known_hosts filename"),
    "ssh_strict_host_key_checking": DataMeta(
        "SSH strict host key checking",
        "accept-new",
    ),
    "ssh_paramiko_connect_kwargs": DataMeta(
        "Override keyword arguments passed into Paramiko's ``SSHClient.connect``"
    ),
    "ssh_connect_retries": DataMeta("Number of tries to connect via ssh", 0),
    "ssh_connect_retry_min_delay": DataMeta(
        "Lower bound for random delay between retries",
        0.1,
    ),
    "ssh_connect_retry_max_delay": DataMeta(
        "Upper bound for random delay between retries",
        0.5,
    ),
    "ssh_file_transfer_protocol": DataMeta(
        "Protocol to use for file transfers. Can be ``sftp`` or ``scp``.",
        "sftp",
    ),
}


class FileTransferClient(Protocol):
    def getfo(self, remote_filename: str, fl: IO) -> Any | None:
        """
        Get a file from the remote host, writing to the provided file-like object.
        """
        ...

    def putfo(self, fl: IO, remote_filename: str) -> Any | None:
        """
        Put a file to the remote host, reading from the provided file-like object.
        """
        ...


class SSHConnector(BaseConnector):
    """
    Connect to hosts over SSH. This is the default connector and all targets default
    to this meaning you do not need to specify it - ie the following two commands
    are identical:

    .. code:: shell

        pyinfra my-host.net ...
        pyinfra @ssh/my-host.net ...
    """

    __examples_doc__ = """
    An inventory file (``inventory.py``) containing a single SSH target with SSH
    forward agent enabled:

    .. code:: python

        hosts = [
            ("my-host.net", {"ssh_forward_agent": True}),
        ]

    Multiple hosts sharing the same SSH username:

    .. code:: python

        hosts = (
            ["my-host-1.net", "my-host-2.net"],
            {"ssh_user": "ssh-user"},
        )

    Multiple hosts with different SSH usernames:

    .. code:: python

        hosts = [
            ("my-host-1.net", {"ssh_user": "ssh-user"}),
            ("my-host-2.net", {"ssh_user": "other-user"}),
        ]
    """

    handles_execution = True

    data_cls = ConnectorData
    data_meta = connector_data_meta
    data: ConnectorData

    client: Optional[SSHClient] = None

    @override
    @staticmethod
    def make_names_data(name):
        yield "@ssh/{0}".format(name), {"ssh_hostname": name}, []

    def make_paramiko_kwargs(self) -> dict[str, Any]:
        kwargs = {
            "allow_agent": False,
            "look_for_keys": False,
            "hostname": self.data["ssh_hostname"] or self.host.name,
            # Overrides of SSH config via pyinfra host data
            "_pyinfra_ssh_forward_agent": self.data["ssh_forward_agent"],
            "_pyinfra_ssh_config_file": self.data["ssh_config_file"],
            "_pyinfra_ssh_known_hosts_file": self.data["ssh_known_hosts_file"],
            "_pyinfra_ssh_strict_host_key_checking": self.data["ssh_strict_host_key_checking"],
            "_pyinfra_ssh_paramiko_connect_kwargs": self.data["ssh_paramiko_connect_kwargs"],
        }

        for key, value in (
            ("username", self.data["ssh_user"]),
            ("port", int(self.data["ssh_port"] or 0)),
            ("timeout", self.state.config.CONNECT_TIMEOUT),
        ):
            if value:
                kwargs[key] = value

        # Password auth (boo!)
        ssh_password = self.data["ssh_password"]
        if ssh_password:
            kwargs["password"] = ssh_password

        # Key auth!
        ssh_key = self.data["ssh_key"]
        if ssh_key:
            kwargs["pkey"] = get_private_key(
                self.state,
                key_filename=ssh_key,
                key_password=self.data["ssh_key_password"],
            )

        # No key or password, so let's have paramiko look for SSH agents and user keys
        # unless disabled by the user.
        else:
            kwargs["allow_agent"] = self.data["ssh_allow_agent"]
            kwargs["look_for_keys"] = self.data["ssh_look_for_keys"]

        return kwargs

    @override
    def connect(self) -> None:
        retries = self.data["ssh_connect_retries"]

        try:
            while True:
                try:
                    return self._connect()
                except (SSHException, gaierror, socket_error, EOFError):
                    if retries == 0:
                        raise
                    retries -= 1
                    min_delay = self.data["ssh_connect_retry_min_delay"]
                    max_delay = self.data["ssh_connect_retry_max_delay"]
                    sleep(uniform(min_delay, max_delay))
        except SSHException as e:
            raise_connect_error(self.host, "SSH error", e)
        except gaierror as e:
            raise_connect_error(self.host, "Could not resolve hostname", e)
        except socket_error as e:
            raise_connect_error(self.host, "Could not connect", e)
        except EOFError as e:
            raise_connect_error(self.host, "EOF error", e)

    def _connect(self) -> None:
        """
        Connect to a single host. Returns the SSH client if successful. Stateless by
        design so can be run in parallel.
        """

        kwargs = self.make_paramiko_kwargs()
        hostname = kwargs.pop("hostname")
        logger.debug("Connecting to: %s (%r)", hostname, kwargs)

        self.client = SSHClient()

        try:
            self.client.connect(hostname, **kwargs)
        except AuthenticationException as e:
            auth_kwargs = {}

            for key, value in kwargs.items():
                if key in ("username", "password"):
                    auth_kwargs[key] = value
                    continue

                if key == "pkey" and value:
                    auth_kwargs["key"] = self.data["ssh_key"]

            auth_args = ", ".join(
                "{0}={1}".format(key, value) for key, value in auth_kwargs.items()
            )

            raise_connect_error(self.host, "Authentication error ({0})".format(auth_args), e)

        except BadHostKeyException as e:
            remove_entry = e.hostname
            port = self.client._ssh_config.get("port", 22)
            if port != 22:
                remove_entry = f"[{e.hostname}]:{port}"

            logger.warning("WARNING: REMOTE HOST IDENTIFICATION HAS CHANGED!")
            logger.warning(
                ("Someone could be eavesdropping on you right now (man-in-the-middle attack)!"),
            )
            logger.warning("If this is expected, you can remove the bad key using:")
            logger.warning(f"    ssh-keygen -R {remove_entry}")

            raise_connect_error(
                self.host,
                "SSH host key error",
                f"Host key for {e.hostname} does not match.",
            )

        except SSHException as e:
            if self._retry_paramiko_agent_keys(hostname, kwargs, e):
                return
            raise

    @override
    def disconnect(self) -> None:
        self.get_file_transfer_connection.cache.clear()

    def _retry_paramiko_agent_keys(
        self,
        hostname: str,
        kwargs: dict[str, Any],
        error: SSHException,
    ) -> bool:
        # Workaround for Paramiko multi-key bug (paramiko/paramiko#1390).
        if "no existing session" not in str(error).lower():
            return False

        if not kwargs.get("allow_agent"):
            return False

        try:
            agent_keys = list(Agent().get_keys())
        except Exception:
            return False

        if not agent_keys:
            return False

        # Skip the first agent key, since Paramiko already attempted it
        attempt_keys = agent_keys[1:] if len(agent_keys) > 1 else agent_keys

        for agent_key in attempt_keys:
            if self.client is not None:
                try:
                    self.client.close()
                except Exception:
                    pass

            self.client = SSHClient()

            single_key_kwargs = dict(kwargs)
            single_key_kwargs["allow_agent"] = False
            single_key_kwargs["pkey"] = agent_key

            try:
                self.client.connect(hostname, **single_key_kwargs)
                return True
            except AuthenticationException:
                continue
            except SSHException as retry_error:
                if "no existing session" in str(retry_error).lower():
                    continue
                raise retry_error

        return False

    @override
    def run_shell_command(
        self,
        command: StringCommand,
        print_output: bool = False,
        print_input: bool = False,
        **arguments: Unpack["ConnectorArguments"],
    ) -> Tuple[bool, CommandOutput]:
        """
        Execute a command on the specified host.

        Args:
            state (``pyinfra.api.State`` obj): state object for this command
            hostname (string): hostname of the target
            command (string): actual command to execute
            sudo (boolean): whether to wrap the command with sudo
            sudo_user (string): user to sudo to
            get_pty (boolean): whether to get a PTY before executing the command
            env (dict): environment variables to set
            timeout (int): timeout for this command to complete before erroring

        Returns:
            tuple: (exit_code, stdout, stderr)
            stdout and stderr are both lists of strings from each buffer.
        """

        _get_pty = arguments.pop("_get_pty", False)
        _timeout = arguments.pop("_timeout", None)
        _stdin = arguments.pop("_stdin", None)
        _success_exit_codes = arguments.pop("_success_exit_codes", None)

        def execute_command() -> Tuple[int, CommandOutput]:
            unix_command = make_unix_command_for_host(self.state, self.host, command, **arguments)
            actual_command = unix_command.get_raw_value()

            logger.debug(
                "Running command on %s: (pty=%s) %s",
                self.host.name,
                _get_pty,
                unix_command,
            )

            if print_input:
                click.echo("{0}>>> {1}".format(self.host.print_prefix, unix_command), err=True)

            # Run it! Get stdout, stderr & the underlying channel
            assert self.client is not None
            stdin_buffer, stdout_buffer, stderr_buffer = self.client.exec_command(
                actual_command,
                get_pty=_get_pty,
            )

            # Write any stdin and then close it
            if _stdin:
                write_stdin(_stdin, stdin_buffer)
            stdin_buffer.close()

            combined_output = read_output_buffers(
                stdout_buffer,
                stderr_buffer,
                timeout=_timeout,
                print_output=print_output,
                print_prefix=self.host.print_prefix,
            )

            logger.debug("Waiting for exit status...")
            exit_status = stdout_buffer.channel.recv_exit_status()
            logger.debug("Command exit status: %i", exit_status)

            return exit_status, combined_output

        return_code, combined_output = execute_command_with_sudo_retry(
            self.host,
            arguments,
            execute_command,
        )

        if _success_exit_codes:
            status = return_code in _success_exit_codes
        else:
            status = return_code == 0

        return status, combined_output

    @memoize
    def get_file_transfer_connection(self) -> FileTransferClient | None:
        assert self.client is not None
        transport = self.client.get_transport()
        assert transport is not None, "No transport"
        try:
            if self.data["ssh_file_transfer_protocol"] == "sftp":
                logger.debug("Using SFTP for file transfer")
                return SFTPClient.from_transport(transport)
            elif self.data["ssh_file_transfer_protocol"] == "scp":
                logger.debug("Using SCP for file transfer")
                return SCPClient(transport)
            else:
                raise ConnectError(
                    "Unsupported file transfer protocol: {0}".format(
                        self.data["ssh_file_transfer_protocol"],
                    ),
                )
        except SSHException as e:
            raise ConnectError(
                (
                    "Unable to establish SFTP connection. Check that the SFTP subsystem "
                    "for the SSH service at {0} is enabled."
                ).format(self.host),
            ) from e

    def _get_file(self, remote_filename: str, filename_or_io: str | IO):
        with get_file_io(filename_or_io, "wb") as file_io:
            sftp = self.get_file_transfer_connection()
            sftp.getfo(remote_filename, file_io)

    @override
    def get_file(
        self,
        remote_filename: str,
        filename_or_io,
        remote_temp_filename=None,
        print_output: bool = False,
        print_input: bool = False,
        **arguments: Unpack["ConnectorArguments"],
    ) -> bool:
        """
        Download a file from the remote host using SFTP. Supports download files
        with sudo by copying to a temporary directory with read permissions,
        downloading and then removing the copy.
        """

        _sudo = arguments.get("_sudo", False)
        _su_user = arguments.get("_su_user", None)

        if _sudo or _su_user:
            # Get temp file location
            temp_file = remote_temp_filename or self.host.get_temp_filename(remote_filename)

            # Copy the file to the tempfile location and add read permissions
            command = StringCommand(
                "cp", remote_filename, temp_file, "&&", "chmod", "+r", temp_file
            )

            copy_status, output = self.run_shell_command(
                command,
                print_output=print_output,
                print_input=print_input,
                **arguments,
            )

            if copy_status is False:
                logger.error("File download copy temp error: {0}".format(output.stderr))
                return False

            try:
                self._get_file(temp_file, filename_or_io)

            # Ensure that, even if we encounter an error, we (attempt to) remove the
            # temporary copy of the file.
            finally:
                remove_status, output = self.run_shell_command(
                    StringCommand("rm", "-f", temp_file),
                    print_output=print_output,
                    print_input=print_input,
                    **arguments,
                )

            if remove_status is False:
                logger.error("File download remove temp error: {0}".format(output.stderr))
                return False

        else:
            self._get_file(remote_filename, filename_or_io)

        if print_output:
            click.echo(
                "{0}file downloaded: {1}".format(self.host.print_prefix, remote_filename),
                err=True,
            )

        return True

    def _put_file(self, filename_or_io, remote_location):
        logger.debug("Attempting upload of %s to %s", filename_or_io, remote_location)

        attempts = 0
        last_e = None

        while attempts < 3:
            try:
                with get_file_io(filename_or_io) as file_io:
                    sftp = self.get_file_transfer_connection()
                    sftp.putfo(file_io, remote_location)
                return
            except OSError as e:
                logger.warning(f"Failed to upload file, retrying: {e}")
                attempts += 1
                last_e = e

        if last_e is not None:
            raise last_e

    @override
    def put_file(
        self,
        filename_or_io,
        remote_filename,
        remote_temp_filename=None,
        print_output: bool = False,
        print_input: bool = False,
        **arguments: Unpack["ConnectorArguments"],
    ) -> bool:
        """
        Upload file-ios to the specified host using SFTP. Supports uploading files
        with sudo by uploading to a temporary directory then moving & chowning.
        """

        original_arguments = arguments.copy()

        _sudo = arguments.pop("_sudo", False)
        _sudo_user = arguments.pop("_sudo_user", False)
        _doas = arguments.pop("_doas", False)
        _doas_user = arguments.pop("_doas_user", False)
        _su_user = arguments.pop("_su_user", None)

        # sudo/su are a little more complicated, as you can only sftp with the SSH
        # user connected, so upload to tmp and copy/chown w/sudo and/or su_user
        if _sudo or _doas or _su_user:
            # Get temp file location
            temp_file = remote_temp_filename or self.host.get_temp_filename(remote_filename)
            self._put_file(filename_or_io, temp_file)

            # Make sure our sudo/su user can access the file
            other_user = _su_user or _sudo_user or _doas_user
            if other_user:
                status, output = self.run_shell_command(
                    StringCommand("setfacl", "-m", f"u:{other_user}:r", temp_file),
                    print_output=print_output,
                    print_input=print_input,
                    **arguments,
                )

                if status is False:
                    logger.error("Error on handover to sudo/su user: {0}".format(output.stderr))
                    return False

            # Execute run_shell_command w/sudo, etc
            command = StringCommand("cp", temp_file, QuoteString(remote_filename))

            status, output = self.run_shell_command(
                command,
                print_output=print_output,
                print_input=print_input,
                **original_arguments,
            )

            if status is False:
                logger.error("File upload error: {0}".format(output.stderr))
                return False

            # Delete the temporary file now that we've successfully copied it
            command = StringCommand("rm", "-f", temp_file)

            status, output = self.run_shell_command(
                command,
                print_output=print_output,
                print_input=print_input,
                **arguments,
            )

            if status is False:
                logger.error("Unable to remove temporary file: {0}".format(output.stderr))
                return False

        # No sudo and no su_user, so just upload it!
        else:
            self._put_file(filename_or_io, remote_filename)

        if print_output:
            click.echo(
                "{0}file uploaded: {1}".format(self.host.print_prefix, remote_filename),
                err=True,
            )

        return True

    @override
    def check_can_rsync(self) -> None:
        if self.data["ssh_key_password"]:
            raise NotImplementedError(
                "Rsync does not currently work with SSH keys needing passwords."
            )

        if self.data["ssh_password"]:
            raise NotImplementedError("Rsync does not currently work with SSH passwords.")

        if not which("rsync"):
            raise NotImplementedError("The `rsync` binary is not available on this system.")

    @override
    def rsync(
        self,
        src: str,
        dest: str,
        flags: Iterable[str],
        print_output: bool = False,
        print_input: bool = False,
        **arguments: Unpack["ConnectorArguments"],
    ):
        _sudo = arguments.pop("_sudo", False)
        _sudo_user = arguments.pop("_sudo_user", False)

        hostname = self.data["ssh_hostname"] or self.host.name
        user = self.data["ssh_user"]
        if user:
            user = "{0}@".format(user)

        ssh_flags = []
        # To avoid asking for interactive input, specify BatchMode=yes
        ssh_flags.append("-o BatchMode=yes")

        known_hosts_file = self.data["ssh_known_hosts_file"]
        if known_hosts_file:
            ssh_flags.append(
                '-o \\"UserKnownHostsFile={0}\\"'.format(shlex.quote(known_hosts_file))
            )  # never trust users

        strict_host_key_checking = self.data["ssh_strict_host_key_checking"]
        if strict_host_key_checking:
            ssh_flags.append(
                '-o \\"StrictHostKeyChecking={0}\\"'.format(shlex.quote(strict_host_key_checking))
            )

        ssh_config_file = self.data["ssh_config_file"]
        if ssh_config_file:
            ssh_flags.append("-F {0}".format(shlex.quote(ssh_config_file)))

        port = self.data["ssh_port"]
        if port:
            ssh_flags.append("-p {0}".format(port))

        ssh_key = self.data["ssh_key"]
        if ssh_key:
            ssh_flags.append("-i {0}".format(ssh_key))

        remote_rsync_command = "rsync"
        if _sudo:
            remote_rsync_command = "sudo rsync"
            if _sudo_user:
                remote_rsync_command = "sudo -u {0} rsync".format(_sudo_user)

        rsync_command = (
            "rsync {rsync_flags} "
            '--rsh "ssh {ssh_flags}" '
            "--rsync-path '{remote_rsync_command}' "
            "{src} {user}{hostname}:{dest}"
        ).format(
            rsync_flags=" ".join(flags),
            ssh_flags=" ".join(ssh_flags),
            remote_rsync_command=remote_rsync_command,
            user=user or "",
            hostname=hostname,
            src=src,
            dest=dest,
        )

        if print_input:
            click.echo("{0}>>> {1}".format(self.host.print_prefix, rsync_command), err=True)

        return_code, output = run_local_process(
            rsync_command,
            print_output=print_output,
            print_prefix=self.host.print_prefix,
        )

        status = return_code == 0
        if not status:
            raise IOError(output.stderr)

        return True
