"""
Connection API implementation using the OpenSSH `ssh(1)` client.

Any SSH options are passed via dictionaries of options, and later translated
to `-o` client CLI options, incl. Hostname, User, Port, IdentityFile, etc.
No "typical" ssh CLI switches are used.

This allows for a nice flexibility from Python code - this module provides
some sensible option defaults (for scripted use), but you are free to
overwrite any options via class or function arguments (where appropriate).

Note that `.cmd()` quotes arguments to really execute individual arguments
as individual arguments in the remote shell, so you need to give it a proper
iterable (like for other Connections), not a single string with spaces.
"""

import logging
import os
import shlex
import subprocess
import tempfile
import threading
import time
from pathlib import Path

from .. import util
from . import Connection

logger = logging.getLogger("atex.connection.ssh")


DEFAULT_OPTIONS = {
    "LogLevel": "ERROR",
    "StrictHostKeyChecking": "no",
    "UserKnownHostsFile": "/dev/null",
    "ConnectionAttempts": "3",
    "ServerAliveCountMax": "4",
    "ServerAliveInterval": "5",
    "TCPKeepAlive": "no",
    "EscapeChar": "none",
    "ExitOnForwardFailure": "yes",
    "RequestTTY": "no",
}


class SSHError(ConnectionError):
    pass


class DisconnectedError(SSHError):
    """
    Raised when an already-connected ssh session goes away (breaks connection).
    """


class NotConnectedError(SSHError):
    """
    Raised when an operation on ssh connection is requested, but the connection
    is not yet open (or has been closed/disconnected).
    """


class ConnectError(SSHError):
    """
    Raised when a to-be-opened ssh connection fails to open.
    """


def _shell_cmd(command, sudo=None):
    """
    Make a command line for running `command` on the target system.
    """
    quoted_args = (shlex.quote(str(arg)) for arg in command)
    if sudo:
        return " ".join((
            "exec", "sudo", "--no-update", "--non-interactive", "--user", sudo, "--", *quoted_args,
        ))
    else:
        return " ".join(("exec", *quoted_args))


def _options_to_cli(options):
    """
    Assemble an `ssh(1)` or `sshpass(1)` command line with `-o` options.
    """
    list_opts = []
    for key, value in options.items():
        if isinstance(value, (list, tuple, set)):
            list_opts += (f"-o{key}={v}" for v in value)
        else:
            list_opts.append(f"-o{key}={value}")
    return list_opts


def _options_to_ssh(options, password=None, extra_cli_flags=()):
    """
    Prefix `sshpass(1)` if password was specified.
    """
    cli_opts = _options_to_cli(options)
    if password:
        return (
            "sshpass", "-p", password,
            "ssh", *extra_cli_flags, "-oBatchMode=no", *cli_opts,
            "ignored_arg",
        )
    else:
        # let cli_opts override BatchMode if specified
        return ("ssh", *extra_cli_flags, *cli_opts, "-oBatchMode=yes", "ignored_arg")


# return a string usable for rsync -e
def _options_to_rsync_e(options, password=None):
    """
    Return a string usable for the rsync `-e` argument.
    """
    cli_opts = _options_to_cli(options)
    batch_mode = "-oBatchMode=no" if password else "-oBatchMode=yes"
    return " ".join(("ssh", *cli_opts, batch_mode))  # no ignored_arg inside -e


def _rsync_host_cmd(*args, options, password=None, sudo=None):
    """
    Assemble a rsync command line, noting that

    - `sshpass` must be before `rsync`, not inside the `-e` argument
    - `ignored_arg` must be passed by user as destination, not inside `-e`
    - `sudo` is part of `--rsync-path`, yet another argument
    """
    return (
        *(("sshpass", "-p", password) if password else ()),
        "rsync",
        "-e", _options_to_rsync_e(options, password=password),
        "--rsync-path", _shell_cmd(("rsync",), sudo=sudo),
        *args,
    )


class StatelessSSHConnection(Connection):
    """
    Implements the Connection API using a `ssh(1)` client using "standalone"
    (stateless) logic - `connect()` and `disconnect()` are no-op, `.cmd()`
    simply executes the ssh client and `.rsync()` executes `rsync -e ssh`.

    Compared to ManagedSSHConnection, this may be slow for many `.cmd()` calls,
    but every call is stateless, there is no persistent connection.

    For only one `.cmd()`, this class is faster than ManagedSSHConnection.
    """

    def __init__(self, options, *, password=None, sudo=None):
        """
        Prepare to connect to an SSH server specified in `options`.

        - If `password` is given, spawn the `ssh` command via `sshpass` and
          pass the password to it.

        - If `sudo` specifies a username, call `sudo(8)` on the remote shell
          to run under a different user on the remote host.
        """
        self.options = DEFAULT_OPTIONS.copy()
        self.options.update(options)
        self.password = password
        self.sudo = sudo
        self._tmpdir = None
        self._master_proc = None

    def connect(self, block=True):
        pass

    def disconnect(self):
        pass

    # have options as kwarg to be compatible with other functions here
    def cmd(self, command, options=None, func=util.subprocess_run, **func_args):
        unified_options = self.options.copy()
        if options:
            unified_options.update(options)
        if command:
            unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
        return func(
            _options_to_ssh(unified_options, password=self.password),
            **func_args,
        )

    def rsync(self, *args, options=None, func=util.subprocess_run, **func_args):
        unified_options = self.options.copy()
        if options:
            unified_options.update(options)
        return func(
            _rsync_host_cmd(
                *args,
                options=unified_options,
                password=self.password,
                sudo=self.sudo,
            ),
            check=True,
            stdin=subprocess.DEVNULL,
            **func_args,
        )


# Note that when ControlMaster goes away (connection breaks), any ssh clients
# connected through it will time out after a combination of
#   ServerAliveCountMax + ServerAliveInterval + ConnectionAttempts
# identical to the ControlMaster process.
# Specifying different values for the clients, to make them exit faster when
# the ControlMaster dies, has no effect. They seem to ignore the options.
#
# If you need to kill the clients quickly after ControlMaster disconnects,
# you need to set up an independent polling logic (ie. every 0.1sec) that
# checks .assert_master() and manually signals the running clients
# when it gets DisconnectedError from it.

class ManagedSSHConnection(Connection):
    """
    Implements the Connection API using one persistently-running `ssh(1)` client
    started in a ControlMaster mode, with additional ssh clients using that
    session to execute remote commands. Similarly, `.rsync()` uses it too.

    This is much faster than StatelessSSHConnection when executing multiple
    commands, but contains a complex internal state (what if ControlMaster
    disconnects?).
    """

    # TODO: thread safety and locking via self.lock ?

    def __init__(self, options, *, password=None, sudo=None):
        """
        Prepare to connect to an SSH server specified in `options`.

        - If `password` is given, spawn the `ssh` command via `sshpass`
        and pass the password to it.

        - If `sudo` specifies a username, call `sudo(8)` on the remote shell
          to run under a different user on the remote host.
        """
        self.lock = threading.RLock()
        self.options = DEFAULT_OPTIONS.copy()
        self.options.update(options)
        self.password = password
        self.sudo = sudo
        self._tmpdir = None
        self._master_proc = None

    def assert_master(self):
        proc = self._master_proc
        if not proc:
            raise NotConnectedError("SSH ControlMaster is not running")
        # we need to consume any potential proc output for the process to
        # actually terminate (stop being a zombie) if it crashes
        out = proc.stdout.read()
        code = proc.poll()
        if code is not None:
            self._master_proc = None
            out = f":\n{out.decode()}" if out else ""
            raise DisconnectedError(
                f"SSH ControlMaster on {self._tmpdir} exited with {code}{out}",
            )

    def disconnect(self):
        proc = self._master_proc
        if not proc:
            return
        logger.info(f"disconnecting: {self.options}")
        proc.kill()
        # don't zombie forever, return EPIPE on any attempts to write to us
        proc.stdout.close()
        proc.wait()
        (self._tmpdir / "control.sock").unlink(missing_ok=True)
        self._master_proc = None

    def connect(self, block=True):
        if not self._tmpdir:
            # _tmpdir_handle just prevents the TemporaryDirectory instance
            # from being garbage collected (and removed on disk)
            # TODO: create/remove it explicitly in connect/disconnect
            #       so the removal happens immediately, even if GC delays cleaning
            self._tmpdir_handle = tempfile.TemporaryDirectory(prefix="atex-ssh-")
            self._tmpdir = Path(self._tmpdir_handle.name)

        sock = self._tmpdir / "control.sock"

        if not self._master_proc:
            logger.info(f"connecting: {self.options}")
            options = self.options.copy()
            options["SessionType"] = "none"
            options["ControlMaster"] = "yes"
            options["ControlPath"] = sock
            self._master_proc = util.subprocess_Popen(
                _options_to_ssh(options, password=self.password),
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                cwd=str(self._tmpdir),
                start_new_session=True,  # resist Ctrl-C
            )
            os.set_blocking(self._master_proc.stdout.fileno(), False)

        # NOTE: ideally, we would .read() before checking .poll() because
        #       if the process writes a lot, it gets stuck in the pipe
        #       (in kernel) and the process never ends; but output-appending
        #       code would be obscure, and ssh(1) never outputs that much ..

        proc = self._master_proc
        if block:
            while proc.poll() is None:
                if sock.exists():
                    break
                time.sleep(0.1)
            else:
                code = proc.poll()
                out = proc.stdout.read()
                self._master_proc = None
                raise ConnectError(
                    f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
                )
        else:
            code = proc.poll()
            if code is not None:
                out = proc.stdout.read()
                self._master_proc = None
                raise ConnectError(
                    f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
                )
            elif not sock.exists():
                raise BlockingIOError("SSH ControlMaster not yet ready")

    def forward(self, forward_type, *spec, cancel=False):
        """
        Add (one or more) ssh forwarding specifications as `spec` to an
        already-connected instance. Each specification has to follow the
        format of LocalForward or RemoteForward (see `ssh_config(5)`).
        Ie. `1234 1.2.3.4:22` or `0.0.0.0:1234 1.2.3.4:22`.

        - `forward_type` must be either LocalForward or RemoteForward.

        - If `cancel` is `True`, cancel the forwarding instead of adding it.
        """
        assert forward_type in ("LocalForward", "RemoteForward")
        self.assert_master()
        options = DEFAULT_OPTIONS.copy()
        options[forward_type] = spec
        options["ControlPath"] = self._tmpdir / "control.sock"
        action = "forward" if not cancel else "cancel"
        util.subprocess_run(
            _options_to_ssh(options, extra_cli_flags=("-O", action)),
            check=True,
        )

    def cmd(self, command, *, options=None, func=util.subprocess_run, **func_args):
        self.assert_master()
        unified_options = self.options.copy()
        if options:
            unified_options.update(options)
        if command:
            unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
        unified_options["ControlPath"] = self._tmpdir / "control.sock"
        return func(
            _options_to_ssh(unified_options),
            **func_args,
        )

    def rsync(self, *args, options=None, func=util.subprocess_run, **func_args):
        self.assert_master()
        unified_options = self.options.copy()
        if options:
            unified_options.update(options)
        unified_options["ControlPath"] = self._tmpdir / "control.sock"
        return func(
            _rsync_host_cmd(
                *args,
                options=unified_options,
                sudo=self.sudo,
            ),
            check=True,
            stdin=subprocess.DEVNULL,
            **func_args,
        )
