#!/usr/bin/env python
# vim:fileencoding=utf-8
# License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>

import os
import re
import subprocess
from typing import Callable, Dict, Iterable, Iterator, Sequence, Tuple

from kitty.complete import Completions, complete_files_and_dirs, debug
from kitty.types import run_once

debug


def lines_from_file(path: str) -> Iterator[str]:
    try:
        f = open(os.path.expanduser(path))
    except OSError:
        pass
    else:
        yield from f


def lines_from_command(*cmd: str) -> Iterator[str]:
    try:
        output = subprocess.check_output(cmd).decode('utf-8')
    except Exception:
        return
    yield from output.splitlines()


def parts_yielder(lines: Iterable[str], pfilter: Callable[[str], Iterator[str]]) -> Iterator[str]:
    for line in lines:
        yield from pfilter(line)


def hosts_from_config_lines(line: str) -> Iterator[str]:
    parts = line.strip().split()
    if len(parts) > 1 and parts[0] == 'Host':
        yield parts[1]


def hosts_from_known_hosts(line: str) -> Iterator[str]:
    parts = line.strip().split()
    if parts:
        yield re.sub(r':\d+$', '', parts[0])


def hosts_from_hosts(line: str) -> Iterator[str]:
    line = line.strip()
    if not line.startswith('#'):
        parts = line.split()
        if parts:
            yield parts[0]
        if len(parts) > 1:
            yield parts[1]
        if len(parts) > 2:
            yield parts[2]


def iter_known_hosts() -> Iterator[str]:
    yield from parts_yielder(lines_from_file('~/.ssh/config'), hosts_from_config_lines)
    yield from parts_yielder(lines_from_file('~/.ssh/known_hosts'), hosts_from_known_hosts)
    yield from parts_yielder(lines_from_file('/etc/ssh/ssh_known_hosts'), hosts_from_known_hosts)
    yield from parts_yielder(lines_from_file('/etc/hosts'), hosts_from_hosts)
    yield from parts_yielder(lines_from_command('getent', 'hosts'), hosts_from_hosts)


@run_once
def known_hosts() -> Tuple[str, ...]:
    return tuple(sorted(filter(lambda x: '*' not in x and '[' not in x, set(iter_known_hosts()))))


@run_once
def ssh_options() -> Dict[str, str]:
    stderr = subprocess.Popen(['ssh'], stderr=subprocess.PIPE).stderr
    assert stderr is not None
    raw = stderr.read().decode('utf-8')
    ans: Dict[str, str] = {}
    pos = 0
    while True:
        pos = raw.find('[', pos)
        if pos < 0:
            break
        num = 1
        epos = pos
        while num > 0:
            epos += 1
            if raw[epos] not in '[]':
                continue
            num += 1 if raw[epos] == '[' else -1
        q = raw[pos+1:epos]
        pos = epos
        if len(q) < 2 or q[0] != '-':
            continue
        if ' ' in q:
            opt, desc = q.split(' ', 1)
            ans[opt[1:]] = desc
        else:
            ans.update(dict.fromkeys(q[1:], ''))
    return ans


# option help {{{
@run_once
def option_help_map() -> Dict[str, str]:
    ans: Dict[str, str] = {}
    lines = '''
-4  -- force ssh to use IPv4 addresses only
-6  -- force ssh to use IPv6 addresses only
-a  -- disable forwarding of authentication agent connection
-A  -- enable forwarding of the authentication agent connection
-B  -- bind to specified interface before attempting to connect
-b  -- specify interface to transmit on
-C  -- compress data
-c  -- select encryption cipher
-D  -- specify a dynamic port forwarding
-E  -- append log output to file instead of stderr
-e  -- set escape character
-f  -- go to background
-F  -- specify alternate config file
-g  -- allow remote hosts to connect to local forwarded ports
-G  -- output configuration and exit
-i  -- select identity file
-I  -- specify smartcard device
-J  -- connect via a jump host
-k  -- disable forwarding of GSSAPI credentials
-K  -- enable GSSAPI-based authentication and forwarding
-L  -- specify local port forwarding
-l  -- specify login name
-M  -- master mode for connection sharing
-m  -- specify mac algorithms
-N  -- don't execute a remote command
-n  -- redirect stdin from /dev/null
-O  -- control an active connection multiplexing master process
-o  -- specify extra options
-p  -- specify port on remote host
-P  -- use non privileged port
-Q  -- query parameters
-q  -- quiet operation
-R  -- specify remote port forwarding
-s  -- invoke subsystem
-S  -- specify location of control socket for connection sharing
-T  -- disable pseudo-tty allocation
-t  -- force pseudo-tty allocation
-V  -- show version number
-v  -- verbose mode (multiple increase verbosity, up to 3)
-W  -- forward standard input and output to host
-w  -- request tunnel device forwarding
-x  -- disable X11 forwarding
-X  -- enable (untrusted) X11 forwarding
-Y  -- enable trusted X11 forwarding
-y  -- send log info via syslog instead of stderr
'''.splitlines()
    for line in lines:
        line = line.strip()
        if line:
            parts = line.split(maxsplit=2)
            ans[parts[0]] = parts[2]
    return ans
# }}}


def complete_choices(ans: Completions, prefix: str, title: str, key: str, comma_separated: bool) -> None:
    choices = {}
    word_transforms = {}
    effective_prefix = prefix
    hidden_prefix = ''
    if comma_separated:
        effective_prefix = prefix.split(',')[-1]
        hidden_prefix = ','.join(prefix.split(',')[:-1])
        if hidden_prefix:
            hidden_prefix += ','
    for line in lines_from_command('ssh', '-Q', key):
        q = line.strip()
        if q.startswith(effective_prefix):
            if comma_separated:
                tq = q
                q = hidden_prefix + q + ','
                word_transforms[q] = tq
            choices[q] = ''
    ans.add_match_group(title, choices, trailing_space=not comma_separated, word_transforms=word_transforms)


def complete_arg(ans: Completions, option_flag: str, prefix: str = '') -> None:
    options = ssh_options()
    option_name = options.get(option_flag[1:])
    if option_name.endswith('file') or option_name.endswith('path'):
        return complete_files_and_dirs(ans, prefix, option_name)
    choices = {
        'mac_spec': ('MAC algorithms', 'mac', True),
        'cipher_spec': ('encryption ciphers', 'cipher', True),
    }
    if option_name in choices:
        return complete_choices(ans, prefix, *choices[option_name])


def complete_destination(ans: Completions, prefix: str = '') -> None:
    result = (k for k in known_hosts() if k.startswith(prefix))
    ans.add_match_group('remote host name', result)


def complete_option(ans: Completions, prefix: str = '-') -> None:
    hm = option_help_map()
    if len(prefix) <= 1:
        result = {k: v for k, v in hm.items() if k.startswith(prefix)}
        ans.add_match_group('option', result)
    else:
        ans.add_match_group('option', {prefix: ''})


def complete(ans: Completions, words: Sequence[str], new_word: bool) -> None:
    options = ssh_options()
    expecting_arg = False
    types = ['' for i in range(len(words))]
    for i, word in enumerate(words):
        if expecting_arg:
            types[i] = 'arg'
            expecting_arg = False
            continue
        if word.startswith('-'):
            types[i] = 'option'
            if len(word) == 2 and options.get(word[1]):
                expecting_arg = True
            continue
        types[i] = 'destination'
        break
    if new_word:
        if words:
            if expecting_arg:
                return complete_arg(ans, words[-1])
        return complete_destination(ans)
    if words:
        if types[-1] == 'arg' and len(words) > 1:
            return complete_arg(ans, words[-2], words[-1])
        if types[-1] == 'destination':
            return complete_destination(ans, words[-1])
        if types[-1] == 'option':
            return complete_option(ans, words[-1])