.dotfile/fish/functions/ff-ssh.fish

# ff ssh: local init (keys + config.d), copy key + add host, copy config to remote

function __ff_ssh_has_key
    for f in $HOME/.ssh/id_*
        if test -f $f; and not string match -q '*.pub' -- $f
            return 0
        end
    end
    return 1
end

function __ff_ssh_ensure_keys
    if __ff_ssh_has_key
        echo "SSH key exists, skip generation"
        return 0
    end

    set -l key_path "$HOME/.ssh/id_ed25519"
    set -l comment "$USER@"(hostname -s 2>/dev/null; or echo localhost)
    mkdir -p "$HOME/.ssh"
    chmod 700 "$HOME/.ssh"

    echo "No SSH key found, generating ed25519: $key_path"
    if test -t 0
        ssh-keygen -t ed25519 -f $key_path -C $comment
    else
        ssh-keygen -t ed25519 -f $key_path -C $comment -N ""
        echo "Non-interactive mode: created with empty passphrase"
    end
end

function __ff_ssh_ensure_config_d
    set -l ssh_config "$HOME/.ssh/config"
    set -l include_line "Include ~/.ssh/config.d/*"
    mkdir -p "$HOME/.ssh/config.d"
    chmod 700 "$HOME/.ssh"

    if not test -f $ssh_config
        printf '%s\n' $include_line > $ssh_config
        chmod 600 $ssh_config
        echo "Created $ssh_config with $include_line"
    else if grep -qE 'Include[[:space:]]+.*config\.d' $ssh_config
        echo "config.d Include already in $ssh_config"
    else
        printf '\n# config.d\n%s\n' $include_line >> $ssh_config
        echo "Appended to $ssh_config: $include_line"
    end
end

function __ff_ssh_resolve_path --argument path
    if test -L $path
        set path (readlink $path)
        if not string match -q '/*' -- $path
            set path "$HOME/.ssh/$path"
        end
    end
    echo $path
end

function __ff_ssh_remote_has_include --argument host
    ssh -q $host 'test -f ~/.ssh/config && grep -qE "Include[[:space:]]+.*config\.d" ~/.ssh/config'
end

function __ff_ssh_remote_ensure_include --argument host
    ssh $host 'mkdir -p ~/.ssh/config.d && chmod 700 ~/.ssh
        if [ -f ~/.ssh/config ] && grep -qE "Include[[:space:]]+.*config\.d" ~/.ssh/config; then
            echo "Remote config.d Include already exists"
        elif [ -f ~/.ssh/config ]; then
            printf "\n# config.d\nInclude ~/.ssh/config.d/*\n" >> ~/.ssh/config
            chmod 600 ~/.ssh/config
            echo "Appended config.d Include on remote"
        else
            printf "%s\n" "Include ~/.ssh/config.d/*" > ~/.ssh/config
            chmod 600 ~/.ssh/config
            echo "Created remote config with config.d Include"
        fi'
end

function __ff_ssh_resolve_copy_src --argument file
    set -l src $file
    if test -f $src
        echo $src
        return 0
    end
    set -l in_config_d "$HOME/.ssh/config.d/$file"
    if test -f $in_config_d
        echo $in_config_d
        return 0
    end
    set -l dotfile_src "$HOME/.dotfile/ssh/config.d/$file"
    if test -f $dotfile_src
        echo $dotfile_src
        return 0
    end
    return 1
end

function __ff_ssh_copy_remote
    set -l host $argv[1]
    set -l files $argv[2..-1]
    set -l ssh_dir "$HOME/.ssh"
    set -l main_config "$ssh_dir/config"

    if test -z "$host"
        echo "Usage: ff ssh copy <host> <config-file>..."
        echo "  host:  SSH alias or user@hostname"
        echo "  file:  path, or name under ~/.ssh/config.d/"
        return 1
    end

    if test (count $files) -eq 0
        echo "Usage: ff ssh copy <host> <config-file>..."
        echo ""
        echo "Examples:"
        echo "  ff ssh copy mynas 30-hanxi.conf"
        echo "  ff ssh copy mynas ~/.dotfile/ssh/config.d/20-personal.conf"
        if test -d "$ssh_dir/config.d"
            echo ""
            echo "Available in ~/.ssh/config.d/:"
            ls -1 "$ssh_dir/config.d/"
        end
        return 1
    end

    echo "Checking remote Include on $host ..."
    ssh -q $host "mkdir -p ~/.ssh/config.d && chmod 700 ~/.ssh"

    if __ff_ssh_remote_has_include $host
        echo "Remote config.d Include: OK"
    else
        echo "Remote missing config.d Include, adding ..."
        __ff_ssh_remote_ensure_include $host
    end

    for f in $files
        set -l src (__ff_ssh_resolve_copy_src $f 2>/dev/null)
        if test $status -ne 0 -o -z "$src"
            echo "Error: file not found: $f"
            return 1
        end

        set -l resolved_main (__ff_ssh_resolve_path $main_config)
        set -l name (basename $src)

        if test "$src" = "$resolved_main"; or test "$name" = config
            set -l ts (date +%Y%m%d%H%M%S)
            ssh -q $host "test -f ~/.ssh/config && cp ~/.ssh/config ~/.ssh/config.bak.$ts || true"
            scp -q $src "$host:.ssh/config"
            ssh -q $host "chmod 600 ~/.ssh/config"
            echo "Copied $src -> $host:~/.ssh/config"
        else
            scp -q $src "$host:.ssh/config.d/$name"
            ssh -q $host "chmod 600 ~/.ssh/config.d/$name"
            echo "Copied $src -> $host:~/.ssh/config.d/$name"
        end
    end

    echo "Done."
end

# ---- ssh-style argument parsing -------------------------------------------

# ssh short options that consume a following value
function __ff_ssh_opt_takes_value --argument-names flag
    contains -- $flag -b -c -D -E -e -F -I -i -J -L -l -m -O -o -p -Q -R -S -W -w
end

# First non-option argument of `ssh`-style args (the [user@]host)
function __ff_ssh_find_target
    set -l n (count $argv)
    set -l i 1
    while test $i -le $n
        set -l a $argv[$i]
        if test "$a" = '--'
            set i (math $i + 1)
            if test $i -le $n
                echo $argv[$i]
            end
            return 0
        end
        if string match -q -- '-*' $a
            set -l flag (string sub -s 1 -l 2 -- $a)
            if test (string length -- $a) -eq 2; and __ff_ssh_opt_takes_value $flag
                set i (math $i + 2)
                continue
            end
            set i (math $i + 1)
            continue
        end
        echo $a
        return 0
    end
    return 1
end

# Value of a short option, supporting both `-i file` and `-ifile`
function __ff_ssh_argv_value --argument-names flag
    set -l args $argv[2..-1]
    set -l n (count $args)
    set -l flen (string length -- $flag)
    set -l i 1
    while test $i -le $n
        set -l a $args[$i]
        if test "$a" = "$flag"
            set -l j (math $i + 1)
            if test $j -le $n
                echo $args[$j]
            end
            return 0
        end
        if test (string length -- $a) -gt $flen; and string match -q -- "$flag*" $a
            string sub -s (math $flen + 1) -- $a
            return 0
        end
        set i (math $i + 1)
    end
    return 1
end

# One value from `ssh -G` output (stdin)
function __ff_ssh_gval --argument-names key
    while read -l line
        set -l trimmed (string trim -- $line)
        set -l k (string split -m1 ' ' -- $trimmed)[1]
        if test (string lower -- $k) = (string lower -- $key)
            string replace -r '^\S+\s+' '' -- $trimmed
            return 0
        end
    end
end

# File that already declares `Host <alias>` (main config or config.d)
function __ff_ssh_find_alias_file --argument-names alias_name
    set -l files $HOME/.ssh/config
    if test -d $HOME/.ssh/config.d
        set files $files $HOME/.ssh/config.d/*
    end
    for f in $files
        test -f $f; or continue
        if awk -v want=(string lower -- $alias_name) '
                tolower($1) == "host" {
                    for (i = 2; i <= NF; i++)
                        if (tolower($i) == want) hit = 1
                }
                END { exit(hit ? 0 : 1) }
            ' $f
            echo $f
            return 0
        end
    end
    return 1
end

# 复制公钥到远端（ssh-copy-id），失败不阻断配置写入
function __ff_ssh_copy_id
    set -l target $argv[1]
    set -l port $argv[2]
    set -l identity $argv[3]
    set -l proxyjump $argv[4]
    set -l proxycommand $argv[5]

    if not type -q ssh-copy-id
        echo "ff ssh: 未找到 ssh-copy-id，跳过公钥复制（可手动执行 ssh-copy-id $target）" >&2
        return 1
    end

    set -l cmd ssh-copy-id
    if test -n "$port"; and test "$port" != 22
        set cmd $cmd -p $port
    end
    if test -n "$identity"
        set cmd $cmd -i $identity
    end
    if test -n "$proxyjump"
        set cmd $cmd -o "ProxyJump=$proxyjump"
    end
    if test -n "$proxycommand"
        set cmd $cmd -o "ProxyCommand=$proxycommand"
    end
    set cmd $cmd $target

    echo "复制公钥: $cmd"
    $cmd
end

# `ff ssh a@b.com` / `ff ssh -p 2222 -i key root@host`
# 确保本地 key -> ssh-copy-id -> 写 ~/.ssh/config.d/<alias>.conf
function __ff_ssh_add_host
    set -l args
    set -l do_copy 1
    for a in $argv
        if test "$a" = --no-copy
            set do_copy 0
        else
            set args $args $a
        end
    end

    set -l target (__ff_ssh_find_target $args)
    if test -z "$target"
        echo "ff ssh: 未找到 [user@]hostname 参数" >&2
        return 1
    end

    # ssh://user@host:port/path -> user / host / port
    set -l plain (string replace -r '^ssh://' '' -- $target)
    set plain (string replace -r '/.*$' '' -- $plain)
    set -l user ""
    if string match -q '*@*' -- $plain
        set user (string replace -r '@[^@]*$' '' -- $plain)
    end
    set -l hostport (string replace -r '^.*@' '' -- $plain)
    set -l port ""
    if string match -q '*:*' -- $hostport
        set port (string replace -r '^.*:' '' -- $hostport)
        set hostport (string replace -r ':[^:]*$' '' -- $hostport)
    end
    set -l host $hostport
    if test -z "$host"
        echo "ff ssh: 无法解析主机名: $target" >&2
        return 1
    end
    test -n "$port"; or set port (__ff_ssh_argv_value -p $args)

    set -l alias_name $host
    if test -t 0
        read -P "SSH 别名 [$host]: " answer
        if test $status -eq 0; and test -n "$answer"
            set alias_name $answer
        end
    end
    if not string match -qr '^[A-Za-z0-9][A-Za-z0-9._-]*$' -- $alias_name
        echo "ff ssh: 非法别名 '$alias_name'（只允许字母、数字、点、下划线、连字符）" >&2
        return 1
    end

    set -l existing (__ff_ssh_find_alias_file $alias_name)
    if test -n "$existing"
        echo "SSH 别名 '$alias_name' 已存在于 $existing，跳过写入"
        return 0
    end

    # 用 ssh -G 解析最终生效的 HostName/User/Port/Proxy*
    set -l resolved (ssh -G $args 2>/dev/null)
    set -l rhost (printf '%s\n' $resolved | __ff_ssh_gval hostname)
    set -l ruser (printf '%s\n' $resolved | __ff_ssh_gval user)
    set -l rport (printf '%s\n' $resolved | __ff_ssh_gval port)
    set -l proxyjump (printf '%s\n' $resolved | __ff_ssh_gval proxyjump)
    set -l proxycommand (printf '%s\n' $resolved | __ff_ssh_gval proxycommand)
    test -n "$rhost"; or set rhost $host
    test -n "$user"; or set user $ruser
    test -n "$port"; or set port $rport
    test -n "$port"; or set port 22
    test "$proxyjump" = none; and set proxyjump ""
    test "$proxycommand" = none; and set proxycommand ""

    # IdentityFile: ssh -G 会列出所有默认密钥，仅在显式指定时记录
    set -l identity ""
    set -l explicit_identity (__ff_ssh_argv_value -i $args)
    if test -n "$explicit_identity"
        set identity $explicit_identity
    else if string match -q '*IdentityFile*' -- $args
        set identity (printf '%s\n' $resolved | __ff_ssh_gval identityfile)
    end
    if string match -q -- "$HOME/*" "$identity"
        set identity "~/"(string replace -r "^$HOME/" '' -- $identity)
    end

    # ssh-copy-id 也要用同一把密钥（-i 或 -o IdentityFile=），否则会复制默认 key
    set -l copy_identity $explicit_identity
    if test -z "$copy_identity"; and test -n "$identity"
        set copy_identity (string replace -r '^~/' "$HOME/" -- $identity)
    end

    # 先复制公钥，这样写完配置就能直接 ssh
    if test "$do_copy" = 1
        __ff_ssh_ensure_keys
        __ff_ssh_copy_id "$target" "$port" "$copy_identity" "$proxyjump" "$proxycommand"
        or echo "ff ssh: 公钥复制失败，已继续写入配置（可重试: ff ssh add $target）" >&2
    end

    set -l dir "$HOME/.ssh/config.d"
    mkdir -p $dir
    chmod 700 "$HOME/.ssh" $dir
    set -l conf_file "$dir/$alias_name.conf"
    if test -f $conf_file
        cp -p $conf_file "$conf_file.backup."(date +%Y%m%d%H%M%S)
    end

    set -l tmpdir /tmp
    test -n "$TMPDIR"; and set tmpdir $TMPDIR
    set -l tmp (mktemp "$tmpdir/ff-ssh.XXXXXX")
    begin
        printf 'Host %s\n' $alias_name
        printf '    HostName %s\n' $rhost
        if test -n "$user"
            printf '    User %s\n' $user
        end
        printf '    Port %s\n' $port
        if test -n "$identity"
            printf '    IdentityFile %s\n' $identity
        end
        if test -n "$proxyjump"
            printf '    ProxyJump %s\n' $proxyjump
        end
        if test -n "$proxycommand"
            printf '    ProxyCommand %s\n' $proxycommand
        end
    end > $tmp
    chmod 600 $tmp
    mv $tmp $conf_file

    echo "已写入 SSH 配置: $conf_file"
    echo "以后可直接运行: ssh $alias_name"
end

function __ff_ssh_usage
    echo "用法: ff ssh [init | <ssh 参数> | add <ssh 参数> | run <ssh 参数> | copy <host> <file>... | help]"
    echo ""
    echo "  (无参数) / init     确保本地 SSH key 存在，并启用 ~/.ssh/config.d Include"
    echo "  add <ssh 参数>      确保本地 key，ssh-copy-id 后写入 ~/.ssh/config.d/<别名>.conf"
    echo "  run <ssh 参数>      同 add，写完后立即连接（此时应已免密）"
    echo "  <ssh 参数>          裸参数等同 add（例: ff ssh -p 2222 root@1.2.3.4）"
    echo "  copy <host> <file>  检查远端 Include，然后 scp 指定配置文件到远端"
    echo "  help                显示此帮助"
    echo ""
    echo "选项:"
    echo "  --no-copy           只写配置，不执行 ssh-copy-id"
    echo ""
    echo "保留子命令: init / copy / add / run / help（同名主机请用 ssh 直接连接）"
    echo ""
    echo "示例:"
    echo "  ff ssh                              # 初始化本地 key 与 config.d"
    echo "  ff ssh a@b.com                      # 复制公钥 + 添加 Host 配置（交互式确认别名）"
    echo "  ff ssh -p 2222 -i ~/.ssh/id_x root@1.2.3.4"
    echo "  ff ssh --no-copy a@b.com            # 只写配置，不复制公钥"
    echo "  ff ssh run a@b.com                  # 复制公钥 + 写配置并立即连接"
    echo "  ff ssh copy mynas 30-hanxi.conf     # 推送配置片段到远端"
end

function ff_ssh
    set -l sub (test (count $argv) -ge 1; and echo $argv[1]; or echo init)

    switch $sub
        case init ''
            __ff_ssh_ensure_keys
            __ff_ssh_ensure_config_d
        case add
            __ff_ssh_ensure_config_d
            __ff_ssh_add_host $argv[2..-1]
        case run
            set -l run_args
            for a in $argv[2..-1]
                test "$a" = --no-copy; and continue
                set run_args $run_args $a
            end
            __ff_ssh_ensure_config_d
            __ff_ssh_add_host $argv[2..-1]; or return 1
            command ssh $run_args
        case copy
            __ff_ssh_copy_remote $argv[2..-1]
        case help -h --help
            __ff_ssh_usage
        case '*'
            # 裸参数与 ssh 原生参数一致：有主机就记录配置，否则透传给 ssh
            if __ff_ssh_find_target $argv >/dev/null
                __ff_ssh_ensure_config_d
                __ff_ssh_add_host $argv
            else
                command ssh $argv
            end
    end
end