#!/usr/bin/env bash
# connect to iot services
THIS_DIR=$( cd "$( dirname "${BASH_SOURCE[0]:-${(%):-%x}}" )" && pwd )
source "$THIS_DIR/../tools/common.sh"
RIOT_TRUST_CLIENT=${RIOT_TRUST_CLIENT:-${DFS_TRUST:-0}}
RIOT_TRUST_SERVER=${RIOT_TRUST_SERVER:-0}
RIOT_EXTRA_OPTIONS=""

# config
RIOT_CONFIG_FILES=(
    "$DOTFILES/riot-config.sh"
    "$HOME/.config/riot-config.sh"
    "riot-config.sh"
)
for file in "${RIOT_CONFIG_FILES[@]}"; do
    if [[ -f "$file" ]]; then
        source "$file"
    fi
done

# check if port number valid
check_port() {
    ( echo $1 | grep -qxE "[1-9][0-9]{0,4}" ) || return 1
    test $1 -lt 65536 -a $1 -gt 0 || return 1
    return 0
}

# check if username valid
check_username() {
    ( echo $1 | grep -qxE "^[a-z][-a-z0-9_]*\$" ) || return 1
    return 0
}

# get single server setting
# may be called more than once
get_server_meta() {
    # returns:
    RET_HOSTNAME=""
    RET_TRUST_SERVER=0
    RET_PORT=""  # optional
    RET_USERNAME=""  # optional
    RET_JUMP_SERVER=""  # optional
    # body
    local remote="$1"
    # if in the form user@...
    if [[ "$remote" == *@* ]]; then
        RET_USERNAME=${remote%%@*}
        remote=${remote#*@}
        check_username $RET_USERNAME || fmt_warning \"$RET_USERNAME\" is not a valid unix username
    fi
    # if in the form ...:22
    if [[ "$remote" == "["*"]":* || ( "$remote" != "["*"]" && "$remote" == *:* ) ]]; then
        RET_PORT=${remote##*:}
        remote=${remote%:*}
        check_port $RET_PORT || fmt_fatal invalid port number \"$RET_PORT\"
    fi
    # presets -- match remote
    local remote_func="$remote.remote"
    if is_function "$remote_func"; then
        "$remote_func"
    fi
    # presets -- match domain
    RET_HOSTNAME=${remote}
    local domain=${remote##*.}
    local host=${remote%.*}
    # if there's no dot
    if [[ "$host" == "$domain" && "$host" != "["*"]" ]]; then
        domain="default"
    fi
    local domain_func="$domain.domain"
    if is_function "$domain_func"; then
        "$domain_func"
    elif is_function ".domain"; then
        ".domain"
    fi
}

parse_remote() {
    # remote setting, including jump servers
    # called for every remote
    # provides:
    SERVER=""
    TRUST_SERVER=1
    PORT=""  # optional
    USERNAME=""  # optional
    SSH_OPTIONS="-o RequestTTY=yes"  # optional
    if [[ "$RIOT_TRUST_CLIENT" == "1" ]]; then
        SSH_OPTIONS="$SSH_OPTIONS -o PermitLocalCommand=yes"
        if [[ "$(get_os_type)" != "msys" ]]; then
            test "$DFS_DRY_RUN" = "1" || mkdir -p ~/.ssh/master-socket
            SSH_OPTIONS="$SSH_OPTIONS -o ControlMaster=auto -o ControlPath=~/.ssh/master-socket/%C"
        fi
    fi
    # handle input
    local remote="$1"
    local jump_servers=""
    # loop for jump servers
    while [[ -n $remote ]]; do
        local server=${remote%%/*}
        remote=${remote#*/}
        get_server_meta "$server"
        if [[ -n "$RET_JUMP_SERVER" ]]; then
            jump_servers="$jump_servers${jump_servers:+,}$RET_JUMP_SERVER"
        fi
        # only if all servers are trusted
        TRUST_SERVER=$((TRUST_SERVER*RET_TRUST_SERVER))
        if [[ "$server" == "$remote" || -z "$remote" ]]; then
            SERVER="$RET_HOSTNAME"
            PORT="$RET_PORT"
            USERNAME="$RET_USERNAME"
            remote=""
        else
            jump_servers="$jump_servers${jump_servers:+,}$RET_USERNAME${RET_USERNAME:+@}$RET_HOSTNAME${RET_PORT:+:}$RET_PORT"
        fi
    done
    # construct cmd
    if [[ "$RIOT_TRUST_SERVER" == "1" || "$TRUST_SERVER" == "1" ]]; then
        SSH_OPTIONS="$SSH_OPTIONS -o ForwardX11=yes -o ForwardAgent=yes"
    fi
    if [[ -n "$jump_servers" ]]; then
        SSH_OPTIONS="$SSH_OPTIONS -o ProxyJump=$jump_servers"
    fi
}

eval_or_echo() {
    local cmd="$@"
    local DO=""
    if [[ "$DFS_DRY_RUN" == "1" ]]; then
        DO=echo
    fi
    if [[ "$RIOT_USE_TMUX" == "1" ]]; then
        if [[ -z "$RIOT_TMUX_SESS" ]]; then
            RIOT_TMUX_SESS=riot-$(date +%s)
            RIOT_TMUX_WIN=0
            $DO tmux new-session -d -s $RIOT_TMUX_SESS bash -l
        else
            RIOT_TMUX_WIN=$((RIOT_TMUX_WIN+1))
            $DO tmux new-window -t $RIOT_TMUX_SESS:$RIOT_TMUX_WIN -d bash -l
        fi
        $DO tmux send-keys -t $RIOT_TMUX_SESS:$RIOT_TMUX_WIN "$cmd" Enter
    else
        $DO $cmd
    fi
}

# ssh series
prepare_ssh_cmd() {
    local ssh_bin="${1:-ssh}"
    if [[ "$ssh_bin" == "scp" || "$ssh_bin" == "sftp" ]]; then
        local port_param='-P'
    else
        local port_param='-p'
    fi
    echo "$ssh_bin ${PORT:+$port_param} $PORT $SSH_OPTIONS $RIOT_EXTRA_OPTIONS $SCP_SRC $USERNAME${USERNAME:+@}$SERVER $SCP_DST ${@:2}"
}

# ssh
run_ssh()
{
    local cmd="$(prepare_ssh_cmd $@)"
    fmt_note "-->" $cmd
    eval_or_echo $cmd
}

# sshl
run_sshl()
{
    local arg="$1"
    if [[ "$arg" != *":"* ]]; then
        # treat as a port number
        arg=localhost:$arg
    fi
    local port=$(get_free_port)

    SSH_OPTIONS="$SSH_OPTIONS -NC -L $port:$arg"
    local cmd="$(prepare_ssh_cmd ssh)"
    fmt_note "-->" $cmd
    fmt_note "  > please access localhost:$port"
    eval_or_echo $cmd
}

# sshd
run_sshd()
{
    local port=$(get_free_port)

    SSH_OPTIONS="$SSH_OPTIONS -NC -D $port"
    local cmd="$(prepare_ssh_cmd ssh)"
    fmt_note "-->" $cmd
    fmt_note "  > please access localhost:$port"
    eval_or_echo $cmd
}

# scp
run_scp() {
    local src="$1"
    local dst="$2"
    local dst_is_remote=1
    # whoever is ./*, it can't be the remote; whoever not exists on local, it's possible the remote.
    # it is suggested to use ./* for local files.
    if [[ "$src" != "./"* && ( "$dst" == "./"* || ( ! -e "$src" && -e "$dst" ) ) ]]; then
        dst_is_remote=0
    fi
    if [[ "$dst_is_remote" == "1" ]]; then
        SCP_SRC="$src"
        SERVER="$SERVER":"$dst"
    else
        SERVER="$SERVER":"$src"
        SCP_DST="$dst"
    fi
    SSH_OPTIONS="$SSH_OPTIONS -r"
    local cmd="$(prepare_ssh_cmd scp)"
    fmt_note "-->" $cmd
    eval_or_echo $cmd
}

# remove host keys
remove_hostkey()
{
    local key
    if [[ -z "$PORT" || "$PORT" == "22" ]]; then
        key=$SERVER
    else
        key="[$SERVER]:$PORT"
    fi
    ssh-keygen -R "$key"
}

# main
print_help()
{
    fmt_info "usage: $0 <service> [command] [options]"
    cat <<EOF
  available commands:
  - ssh (default)
  - tmux (run ssh in multiple tmux windows)
  - sshl (ssh -L)
  - sshd (ssh -D)
  - zssh
  - sftp
  - rm (remove host keys)
  - ping/ping6 (let the remote to ping the given target)
EOF
}

router() {
    if [[ -z "$1" || "$1" == "-h" || "$1" == "--help" ]]; then
        print_help
        exit
    fi

    while [[ "$1" == -* ]]; do
        if [[ "$1" == "--tmux" ]]; then
            RIOT_USE_TMUX=1
            shift
            continue
        elif [[ "$1" == "-t" || "$1" == "--trust" ]]; then
            RIOT_TRUST_SERVER=1
            shift
            continue
        fi
        RIOT_EXTRA_OPTIONS="$RIOT_EXTRA_OPTIONS $1"
        if [[ "$1" == "-o" ]]; then
            RIOT_EXTRA_OPTIONS="$RIOT_EXTRA_OPTIONS $2"
            shift
        fi
        shift
    done

    IFS=',' read -ra remotes <<< "$1"
    for ((i=0; i<${#remotes[@]}; i++)); do
        remote="${remotes[i]}"
        if [[ -z "$remote" ]]; then
            continue
        fi
        local batch_func="${remote}.batch"
        if is_function "$batch_func"; then
            "$batch_func"
            continue
        fi
        parse_remote "$remote"
        case $2 in
            ssh|tmux|"" )
                [[ "$2" == tmux ]] && RIOT_USE_TMUX=1
                run_ssh ssh "${@:3}"
                ;;
            ping|ping6 )
                run_ssh ssh "${@:2}"
                ;;
            zssh )
                run_ssh zssh
                ;;
            sftp )
                run_ssh sftp
                ;;
            sshl )
                test -n "$3" || fmt_fatal "no target address provided"
                run_sshl "$3"
                ;;
            sshd )
                run_sshd
                ;;
            scp )
                test -n "$3" || fmt_fatal "no source path specified"
                test -n "$4" || fmt_fatal "no destination path specified"
                run_scp "$3" "$4"
                ;;
            rm )
                remove_hostkey
                ;;
            * )
                print_help
                fmt_fatal "unknown command: $2"
                ;;
        esac
    done

    if [[ -n "$RIOT_TMUX_SESS" ]]; then
        tmux attach-session -t $RIOT_TMUX_SESS
    fi
}

router "${GOT_OPTS[@]}"