from __future__ import annotations

__kupfer_name__ = _("SSH Hosts")
__description__ = _("Adds the SSH hosts found in ~/.ssh/config.")
__version__ = "2023-11-04"
__author__ = "Fabian Carlström, Karol Będkowski"

__kupfer_sources__ = ("SSHSource",)
__kupfer_actions__ = ("SSHConnect",)

import os
import typing as ty
from contextlib import suppress

from kupfer import icons, launch
from kupfer.obj import Action
from kupfer.obj.grouping import ToplevelGroupingSource
from kupfer.obj.helplib import FilesystemWatchMixin
from kupfer.obj.hosts import (
    HOST_ADDRESS_KEY,
    HOST_NAME_KEY,
    HOST_SERVICE_NAME_KEY,
    HOST_SERVICE_USER_KEY,
    HostLeaf,
)

if ty.TYPE_CHECKING:
    from gettext import gettext as _


class SSHLeaf(HostLeaf):
    """The SSH host. Using name defined in `Host` or by `Match` statement.
    If `Hostname` is not defined in ssh config - is set to  `name`.
    """

    def __init__(
        self, name: str, hostname: str | None = None, user: str | None = None
    ) -> None:
        slots = {
            HOST_NAME_KEY: name,
            HOST_ADDRESS_KEY: hostname or name,
            HOST_SERVICE_NAME_KEY: "ssh",
        }
        if user:
            slots[HOST_SERVICE_USER_KEY] = user

        HostLeaf.__init__(self, slots, name)

    def get_description(self) -> str:
        host = str(self[HOST_ADDRESS_KEY])
        with suppress(KeyError):
            if user := self[HOST_SERVICE_USER_KEY]:
                host = f"{user}@{host}"

        return _("SSH host: %s") % host

    def get_gicon(self):
        return icons.ComposedIconSmall(self.get_icon_name(), "terminal")


class SSHConnect(Action):
    """Used to launch a terminal connecting to the specified SSH host.
    HOST_NAME_KEY is used for making connections.
    """

    def __init__(self):
        Action.__init__(self, name=_("Connect"))

    def activate(self, leaf, iobj=None, ctx=None):
        host = leaf[HOST_NAME_KEY]
        with suppress(KeyError):
            if user := leaf[HOST_SERVICE_USER_KEY]:
                host = f"{user}@{host}"

        launch.spawn_in_terminal(["ssh", host])

    def get_description(self):
        return _("Connect to SSH host")

    def get_icon_name(self):
        return "network-server"

    def item_types(self):
        yield HostLeaf

    def valid_for_item(self, leaf):
        if leaf.check_key(HOST_SERVICE_NAME_KEY):
            return leaf[HOST_SERVICE_NAME_KEY] == "ssh"

        return False


def _parse_host_stms(args: list[str]) -> ty.Iterator[tuple[str, str | None]]:
    """Load names from `Host` statement. Ignore names with `*` and exclusions."""
    for host in args:
        if "*" in host or host[0] == "!":
            continue

        user, _, host = host.partition("@")
        if user and host:
            yield host, user
        else:
            yield user, None


def _parse_match_stmt(args: list[str]) -> tuple[str, str | None] | None:
    """Parse `Match` arguments and return tuple of (host, optional user)
    or None.

    Load only `User` and `Host` keys from this statement.
    """
    user = None
    host = None
    for key, val in zip(args, args[1:]):
        if "*" in val or val[0] == "!":
            continue

        key = key.lower()
        if key == "host":
            host = val
        elif key == "user":
            user = val

    if not host:
        return None

    return host, user


class SSHSource(ToplevelGroupingSource, FilesystemWatchMixin):
    """Reads ~/.ssh/config and creates leaves for the hosts found."""

    source_scan_interval: int = 3600

    _ssh_home = os.path.expanduser("~/.ssh")
    _config_path = os.path.join(_ssh_home, "config")

    def __init__(self, name=_("SSH Hosts")):
        ToplevelGroupingSource.__init__(self, name, "hosts")
        self._version = 2
        self.monitor_token = None

    def initialize(self):
        ToplevelGroupingSource.initialize(self)
        self.monitor_token = self.monitor_files(self._config_path)

    def monitor_include_file(self, gfile):
        return gfile.get_path() in (self._config_path, self._ssh_home)

    def _get_items(self) -> ty.Iterable[SSHLeaf]:
        with open(self._config_path, encoding="UTF-8") as cfile:
            current_hosts: list[tuple[str, str | None]] = []
            current_hostname: str | None = None
            for line in cfile:
                line = line.strip()
                if not line:
                    continue

                head, *args = line.split()
                head = head.lower()
                if head in ("host", "match"):
                    # new restriction, flush current data
                    for host, user in current_hosts:
                        yield SSHLeaf(host, current_hostname, user)

                    current_hosts.clear()
                    current_hostname = None

                    if head == "host":
                        # process only 'host' restriction; skip wildcard and
                        # negative entries
                        current_hosts = list(_parse_host_stms(args))
                    elif len(args) > 1:
                        # process "key val" parameters of match
                        if hostuser := _parse_match_stmt(args):
                            current_hosts = [hostuser]

                elif head == "hostname" and args:
                    # if found hostname use is as HOST_ADDRESS_KEY
                    current_hostname = args[0]

            for host, user in current_hosts:
                yield SSHLeaf(host, current_hostname, user)

    def get_items(self):
        try:
            return list(self._get_items())
        except OSError as exc:
            self.output_error(exc)
        except UnicodeError as exc:
            self.output_error(
                f"File {self._config_path} not in expected encoding (UTF-8)"
            )
            self.output_error(exc)

        return ()

    def get_description(self):
        return _("SSH hosts as specified in ~/.ssh/config")

    def get_icon_name(self):
        return "network-workgroup"

    def provides(self):
        yield SSHLeaf
