#!/bin/bash

#==============================================================================
#
#          FILE:  sshm
#
#         USAGE:  sshm [alias | ID]
#
#   DESCRIPTION:  A simple SSH connection manager written in Bash.
#                 - 'sshm' enters the management interface.
#                 - 'sshm <alias | ID>' connects to the specified host.
#
#       AUTHOR:  Gemini
#      VERSION:  1.1 (ID connection bug fixed)
#      CREATED:  2025-09-15
#
#==============================================================================

# 配置文件路径
CONFIG_FILE="$HOME/.sshm_hosts"

# 初始化配置文件（如果不存在）
initialize() {
    if [ ! -f "$CONFIG_FILE" ]; then
        touch "$CONFIG_FILE"
        echo ">> Welcome! Configuration file created at ~/.sshm_hosts"
    fi
}

# 打印漂亮的标题
print_header() {
    echo "========================================"
    echo "  SSH Host Manager (sshm) - Bash Edition"
    echo "========================================"
}

# 列出所有已保存的主机
list_hosts() {
    print_header
    echo "Saved SSH Hosts:"
    if [ ! -s "$CONFIG_FILE" ]; then
        echo "  No hosts found. Use 'add' to create one."
    else
        # 使用 awk 美化输出
        awk -F',' 'BEGIN { printf "%-5s %-15s %-25s %-7s %s\n", "ID", "Alias", "User@Host", "Port", "IdentityFile" }
        {
            user_host = $2 "@" $3;
            identity = ($5 == "") ? "default" : $5;
            printf "%-5s %-15s %-25s %-7s %s\n", NR, $1, user_host, $4, identity
        }' "$CONFIG_FILE"
    fi
    echo "----------------------------------------"
}

# 添加一个新的主机
add_host() {
    echo ">> Adding a new host..."
    read -p "Enter Alias: " alias
    if grep -q "^${alias}," "$CONFIG_FILE"; then
        echo "Error: Alias '${alias}' already exists. Please choose another one."
        return
    fi

    read -p "Enter User: " user
    read -p "Enter Hostname/IP: " hostname
    read -p "Enter Port (default 22): " port
    port=${port:-22} # 如果为空，则使用默认值22
    read -p "Enter Identity File path (optional, e.g., ~/.ssh/id_rsa): " identity_file

    echo "${alias},${user},${hostname},${port},${identity_file}" >> "$CONFIG_FILE"
    echo ">> Host '${alias}' added successfully!"
}

# 删除一个主机
delete_host() {
    list_hosts
    if [ ! -s "$CONFIG_FILE" ]; then
        return
    fi
    read -p "Enter the ID or Alias of the host to delete: " identifier

    local temp_file=$(mktemp)
    local deleted=false

    # 判断输入是数字ID还是别名
    if [[ "$identifier" =~ ^[0-9]+$ ]]; then
        # 按行号删除
        awk -v id="$identifier" 'NR != id' "$CONFIG_FILE" > "$temp_file"
        if [ $(wc -l < "$CONFIG_FILE") -ne $(wc -l < "$temp_file") ]; then
            deleted=true
        fi
    else
        # 按别名删除
        grep -v "^${identifier}," "$CONFIG_FILE" > "$temp_file"
        if [ $(wc -l < "$CONFIG_FILE") -ne $(wc -l < "$temp_file") ]; then
            deleted=true
        fi
    fi

    if [ "$deleted" = true ]; then
        mv "$temp_file" "$CONFIG_FILE"
        echo ">> Host '${identifier}' deleted successfully."
    else
        echo "Error: Host with ID or Alias '${identifier}' not found."
        rm "$temp_file"
    fi
}

# 连接到主机 (*** 已修正此函数 ***)
connect_to_host() {
    local identifier="$1"
    local host_info=""

    # 检查输入是数字ID还是别名
    if [[ "$identifier" =~ ^[0-9]+$ ]]; then
        # 如果是数字，使用 awk 按行号获取配置
        host_info=$(awk -v id="$identifier" 'NR == id' "$CONFIG_FILE")
    else
        # 如果不是数字，按别名 grep 获取配置
        host_info=$(grep "^${identifier}," "$CONFIG_FILE")
    fi

    if [ -z "$host_info" ]; then
        echo "Error: Host with ID or Alias '${identifier}' not found."
        return 1 # 使用 return 1 而不是 exit 1，以免退出管理界面
    fi

    # 解析主机信息
    local user=$(echo "$host_info" | cut -d',' -f2)
    local host=$(echo "$host_info" | cut -d',' -f3)
    local port=$(echo "$host_info" | cut -d',' -f4)
    local identity_file=$(echo "$host_info" | cut -d',' -f5)

    local ssh_command="ssh ${user}@${host} -p ${port}"
    if [ -n "$identity_file" ]; then
        # 使用波浪号扩展
        identity_file_expanded=$(eval echo "$identity_file")
        ssh_command+=" -i ${identity_file_expanded}"
    fi

    echo ">> Connecting to ${user}@${host} on port ${port}..."
    eval "$ssh_command"
}


# 显示管理页面的帮助信息
show_help() {
    print_header
    echo "Management Commands:"
    echo "  list      - List all saved SSH hosts."
    echo "  add       - Add a new SSH host."
    echo "  del       - Delete an SSH host by ID or Alias."
    echo "  conn      - Connect to a host by its ID or Alias."
    echo "  help      - Show this help message."
    echo "  exit      - Exit the management interface."
    echo "----------------------------------------"
}


# 管理页面主循环
management_interface() {
    show_help
    while true; do
        read -p "sshm> " cmd args
        case "$cmd" in
            list|ls)
                list_hosts
                ;;
            add)
                add_host
                ;;
            del|rm|delete)
                delete_host
                ;;
            conn|connect)
                local target="$args"
                if [ -z "$target" ]; then
                    read -p "Enter the ID or Alias to connect to: " target
                fi
                if [ -n "$target" ]; then
                    connect_to_host "$target"
                fi
                ;;
            help)
                show_help
                ;;
            exit|quit|q)
                echo ">> Bye!"
                break
                ;;
            "")
                # 允许用户按回车键刷新
                ;;
            *)
                echo "Unknown command: '$cmd'. Type 'help' for a list of commands."
                ;;
        esac
    done
}


# ==================
#   脚本主入口
# ==================

initialize

# 如果没有参数，进入管理模式
if [ $# -eq 0 ]; then
    management_interface
# 如果有一个参数，尝试直接连接
elif [ $# -eq 1 ]; then
    connect_to_host "$1"
# 其他情况，显示用法
else
    echo "Usage: sshm [alias | ID]"
    echo "       Run 'sshm' without arguments to enter the management interface."
    exit 1
fi