#!/bin/bash

# 简单实用的 Docker 容器管理脚本 (优化版)

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m'

# 检查 Docker 及权限
check_docker() {
    if ! command -v docker &> /dev/null; then
        echo -e "${RED}错误: 未检测到 Docker，请先安装。${NC}"
        exit 1
    fi
    # 检查是否有权限执行 docker 命令
    if ! docker ps &> /dev/null; then
        echo -e "${RED}错误: 当前用户无权运行 Docker 命令。${NC}"
        echo -e "${YELLOW}请尝试使用 sudo 运行此脚本，或将用户加入 docker 用户组。${NC}"
        exit 1
    fi
}

# 格式化输出表头
print_table_header() {
    printf "${BLUE}%-5s | %-25s | %-12s | %-20s${NC}\n" "编号" "容器名称" "状态" "镜像"
    echo "------|---------------------------|--------------|----------------------"
}

# 格式化输出行
print_table_row() {
    printf "%-5s | %-25s | %-12b | %-20s\n" "$1" "$2" "$3" "$4"
}

# 显示容器列表
show_containers() {
    echo -e "\n${BLUE}=== 容器列表 ===${NC}"
    print_table_header
    
    local count=0
    # 获取数据并暂存，避免管道导致的子shell变量丢失问题
    mapfile -t lines < <(docker ps -a --format "table {{.Names}}\t{{.Status}}\t{{.Image}}" | tail -n +2)

    for line in "${lines[@]}"; do
        if [ -n "$line" ]; then
            count=$((count + 1))
            name=$(echo "$line" | awk '{print $1}')
            status_raw=$(echo "$line" | awk '{print $2}')
            image=$(echo "$line" | awk '{for(i=3;i<=NF;i++) printf $i" "; print ""}' | sed 's/ $//')
            
            # 截断过长的名称以保持表格整洁
            display_name=$name
            if [ ${#display_name} -gt 23 ]; then display_name="${display_name:0:20}..."; fi
            
            # 截断过长的镜像名
            display_image=$image
            if [ ${#display_image} -gt 18 ]; then display_image="${display_image:0:15}..."; fi

            # 状态美化
            if [[ "$status_raw" == "Up"* ]]; then
                status_display="${GREEN}运行中${NC}"
            elif [[ "$status_raw" == "Exited"* ]]; then
                status_display="${RED}已停止${NC}"
            else
                status_display="${YELLOW}$status_raw${NC}"
            fi
            
            print_table_row "$count" "$display_name" "$status_display" "$display_image"
        fi
    done
    echo
}

# 显示镜像列表 (优化对齐)
show_images() {
    echo -e "\n${PURPLE}=== 镜像列表 ===${NC}"
    printf "${PURPLE}%-5s | %-30s | %-12s | %-8s | %-15s${NC}\n" "编号" "镜像名称:标签" "镜像ID" "大小" "创建时间"
    echo "------|--------------------------------|--------------|----------|----------------"
    
    local count=0
    mapfile -t lines < <(docker images --format "table {{.Repository}}:{{.Tag}}\t{{.ID}}\t{{.Size}}\t{{.CreatedSince}}" | tail -n +2)
    
    for line in "${lines[@]}"; do
        if [ -n "$line" ]; then
            count=$((count + 1))
            repo_tag=$(echo "$line" | awk '{print $1}')
            image_id=$(echo "$line" | awk '{print $2}')
            size=$(echo "$line" | awk '{print $3}')
            created=$(echo "$line" | awk '{for(i=4;i<=NF;i++) printf $i" "; print ""}' | sed 's/ $//')
            
            # 截断处理
            if [ ${#repo_tag} -gt 28 ]; then repo_tag="${repo_tag:0:25}..."; fi
            
            printf "%-5s | %-30s | %-12s | %-8s | %-15s\n" "$count" "$repo_tag" "$image_id" "$size" "$created"
        fi
    done
    echo
}

# 通用获取名称函数 (复用逻辑)
get_name_by_index() {
    local number=$1
    local type=$2 # "container" or "image"
    local cmd=""
    
    if [ "$type" == "container" ]; then
        cmd="docker ps -a --format '{{.Names}}'"
    else
        cmd="docker images --format '{{.ID}}'"
    fi
    
    # 动态获取第N行
    eval "$cmd" | sed -n "${number}p"
}

# 进入容器终端 (新增功能)
enter_container() {
    local container_name=$1
    echo -e "${YELLOW}尝试进入容器: $container_name${NC}"
    echo -e "提示: 输入 ${RED}exit${NC} 可退出容器终端"
    echo
    
    # 优先尝试 bash，失败则尝试 sh
    if docker exec -it "$container_name" /bin/bash 2>/dev/null; then
        return
    elif docker exec -it "$container_name" /bin/sh 2>/dev/null; then
        return
    else
        echo -e "${RED}无法进入容器终端 (可能容器未包含 bash 或 sh)${NC}"
        read -p "按回车键继续..."
    fi
}

# 批量删除容器
batch_delete_containers() {
    local numbers="$1"
    local confirm="$2"
    numbers=$(echo "$numbers" | sed 's/,/ /g')
    local success_count=0
    
    for number in $numbers; do
        if [[ "$number" =~ ^[0-9]+$ ]]; then
            container_name=$(get_name_by_index $number "container")
            if [ -n "$container_name" ]; then
                if [ "$confirm" != "y" ]; then
                    read -p "确认删除容器 $container_name? (y/N): " c
                    [[ "$c" != "y" && "$c" != "Y" ]] && continue
                fi
                
                if docker rm -f "$container_name" &> /dev/null; then
                    echo -e "${GREEN}✓ 删除成功: $container_name${NC}"
                    ((success_count++))
                else
                    echo -e "${RED}✗ 删除失败: $container_name${NC}"
                fi
            else
                echo -e "${RED}✗ 编号 $number 不存在${NC}"
            fi
        fi
    done
    [ $success_count -gt 0 ] && echo -e "${GREEN}共删除 $success_count 个容器${NC}"
}

# 批量删除镜像
batch_delete_images() {
    local numbers="$1"
    local confirm="$2"
    numbers=$(echo "$numbers" | sed 's/,/ /g')
    local success_count=0
    
    for number in $numbers; do
        if [[ "$number" =~ ^[0-9]+$ ]]; then
            image_id=$(get_name_by_index $number "image")
            if [ -n "$image_id" ]; then
                # 获取友好名称用于显示
                image_name=$(docker images --format "{{.Repository}}:{{.Tag}}" | grep -v "REPOSITORY" | grep "$image_id" | head -n 1)
                [ -z "$image_name" ] && image_name=$image_id
                
                if [ "$confirm" != "y" ]; then
                    read -p "确认删除镜像 $image_name? (y/N): " c
                    [[ "$c" != "y" && "$c" != "Y" ]] && continue
                fi
                
                if docker rmi "$image_id" &> /dev/null; then
                    echo -e "${GREEN}✓ 删除成功: $image_name${NC}"
                    ((success_count++))
                else
                    echo -e "${RED}✗ 删除失败 (可能被容器占用): $image_name${NC}"
                fi
            fi
        fi
    done
}

# 清理悬空镜像
clean_dangling_images() {
    local count=$(docker images -f "dangling=true" -q | wc -l)
    if [ "$count" -gt 0 ]; then
        read -p "发现 $count 个悬空镜像，是否清理？(y/N): " confirm
        if [[ "$confirm" == "y" || "$confirm" == "Y" ]]; then
            docker image prune -f
            echo -e "${GREEN}清理完成${NC}"
        fi
    else
        echo -e "${GREEN}系统很干净，没有悬空镜像。${NC}"
    fi
    read -p "按回车键继续..."
}

# 镜像管理菜单
image_management() {
    while true; do
        clear
        echo -e "${PURPLE}========================================${NC}"
        echo -e "${PURPLE}           Docker 镜像管理${NC}"
        echo -e "${PURPLE}========================================${NC}"
        show_images
        
        echo -e "${PURPLE}操作选项:${NC}"
        echo "1. 删除镜像 (例如: 1 或 1,2,3)"
        echo "2. 清空所有镜像 (慎用)"
        echo "3. 清理悬空镜像 (<none>)"
        echo "4. 返回主菜单"
        echo
        read -p "请选择: " choice
        
        case $choice in
            1)
                read -p "请输入镜像编号: " input
                batch_delete_images "$input" "n"
                read -p "按回车键继续..."
                ;;
            2)
                read -p "警告：确定删除所有镜像吗？(yes/N): " confirm
                if [ "$confirm" == "yes" ]; then
                    docker rmi -f $(docker images -q) 2>/dev/null
                    echo -e "${GREEN}已清空所有镜像${NC}"
                fi
                read -p "按回车键继续..."
                ;;
            3) clean_dangling_images ;;
            4) return ;;
            *) echo -e "${RED}无效输入${NC}"; sleep 1 ;;
        esac
    done
}

# 主菜单
main_menu() {
    while true; do
        clear
        echo -e "${BLUE}========================================${NC}"
        echo -e "${BLUE}       Docker 容器可视化管理 v2.0${NC}"
        echo -e "${BLUE}========================================${NC}"
        show_containers
        
        echo -e "${BLUE}容器操作:${NC}                          ${CYAN}高级功能:${NC}"
        printf "%-35s %-35s\n" "1. 启动容器" "7. 进入容器终端 (Exec)"
        printf "%-35s %-35s\n" "2. 停止容器" "8. 实时资源监控 (Stats)"
        printf "%-35s %-35s\n" "3. 重启容器" "9. 镜像管理菜单"
        printf "%-35s\n" "4. 查看日志" 
        echo
        echo -e "${RED}危险操作:${NC}"
        echo "5. 删除容器 (例如: 1 或 1,2,3)"
        echo "6. 清理所有已停止容器"
        echo "0. 退出"
        echo
        
        read -p "请选择操作 [0-9]: " choice
        
        # 处理需要选择容器的操作
        if [[ "1 2 3 4 7" =~ "$choice" ]]; then
            read -p "请输入容器编号: " num
            if [[ ! "$num" =~ ^[0-9]+$ ]]; then
                echo -e "${RED}请输入有效数字${NC}"; sleep 1; continue
            fi
            
            name=$(get_name_by_index $num "container")
            if [ -z "$name" ]; then
                echo -e "${RED}编号不存在${NC}"; sleep 1; continue
            fi

            case $choice in
                1) docker start "$name" && echo -e "${GREEN}已启动 $name${NC}" ;;
                2) docker stop "$name" && echo -e "${GREEN}已停止 $name${NC}" ;;
                3) docker restart "$name" && echo -e "${GREEN}已重启 $name${NC}" ;;
                4) docker logs -f --tail 100 "$name" ;;
                7) enter_container "$name" ;;
            esac
            [ "$choice" != "7" ] && [ "$choice" != "4" ] && sleep 1
            continue
        fi

        # 处理其他操作
        case $choice in
            5)
                read -p "请输入要删除的容器编号 (支持批量): " nums
                batch_delete_containers "$nums" "n"
                read -p "按回车键继续..."
                ;;
            6)
                read -p "确认清理所有已停止的容器？(y/N): " c
                if [[ "$c" == "y" || "$c" == "Y" ]]; then
                    docker container prune -f
                    echo -e "${GREEN}清理完成${NC}"
                fi
                sleep 1
                ;;
            8)
                echo -e "${CYAN}正在打开资源监控 (按 Ctrl+C 退出)...${NC}"
                docker stats
                ;;
            9) image_management ;;
            0) echo -e "${GREEN}再见!${NC}"; exit 0 ;;
            *) echo -e "${RED}无效选择${NC}"; sleep 1 ;;
        esac
    done
}

# 启动
check_docker
main_menu