#!/bin/bash

# set -x
# 脚本作用: 
# 1. bash launch.sh 直接进入容器
# 2. bash launch.sh d / bash launch.sh destroy 删除容器
# 3. bash launch.sh p / bash launch.sh print 打印容器名称
# 4. bash launch.sh -v 启动容器并显示连接命令

# ==================== 配置区域 ====================
# 镜像配置
docker_image="dcu_ai:v1.1"
container_name="yolov5_test_dtk24.04.1"

# 挂载卷配置 (格式: "宿主机路径:容器路径")
# 可以添加多个挂载，用空格分隔
mounts=(
    "/opt/hyhal:/opt/hyhal"
    "/home/wkx/yolov5/datasets:/datasets"
)

# 环境变量配置
environment_vars=(
    "https_proxy=http://10.17.26.123:8086"
    "http_proxy=http://10.17.26.123:8086"
)

# Docker 设备配置
devices=(
    "/dev/kfd"
    "/dev/mkfd"
    "/dev/dri"
)

# Docker 运行参数
shm_size="32G"
extra_args=(
    "--ipc=host"
    "--group-add video"
    "--cap-add=SYS_PTRACE"
    "--security-opt seccomp=unconfined"
    "--network=host"
    "--privileged"
)

# ==================== 配置区域结束 ====================

export CONTAINER_NAME="${container_name}"

# 构建 docker run 命令
build_docker_command() {
    local cmd="docker run -id --name ${CONTAINER_NAME} --shm-size=${shm_size}"
    
    # 添加额外参数
    for arg in "${extra_args[@]}"; do
        cmd="${cmd} ${arg}"
    done
    
    # 添加环境变量
    for env in "${environment_vars[@]}"; do
        cmd="${cmd} -e ${env}"
    done
    
    # 添加设备
    for device in "${devices[@]}"; do
        cmd="${cmd} --device=${device}"
    done
    
    # 添加挂载卷
    for mount in "${mounts[@]}"; do
        cmd="${cmd} -v ${mount}"
    done
    
    cmd="${cmd} ${docker_image} /bin/bash"
    echo "${cmd}"
}

command=$(build_docker_command)

# 检查 CONTAINER_NAME 环境变量是否设置
if [ -z "${CONTAINER_NAME}" ]; then
    echo "环境变量 CONTAINER_NAME 未设置，请设置后再试。"
    exit 1
fi

check_command_result() {
    local exit_status=$1
    
    if [ "$exit_status" -eq 0 ]; then
        echo "命令成功执行"
    else
        echo "命令执行失败，退出状态为: $exit_status"
        exit $exit_status
    fi
}

# 显示容器连接信息
show_container_info() {
    echo ""
    echo "========================================="
    echo "容器名称: ${CONTAINER_NAME}"
    echo "连接命令:"
    echo "docker exec -it ${CONTAINER_NAME} /bin/bash"
    echo "========================================="
    echo ""
}

# 使用 docker ps -a | grep 来查找容器，并计算结果行数
num=$(docker ps -a | grep "$CONTAINER_NAME" | wc -l)

# 检查是否有 -v 参数
show_connect_cmd=false
if [ "$#" -gt 0 ]; then
    for arg in "$@"; do
        if [ "$arg" = "-v" ]; then
            show_connect_cmd=true
            # 移除 -v 参数，避免干扰其他逻辑
            set -- "${@/$arg/}"
            # 重新构建参数列表，去除空参数
            args=()
            for a in "$@"; do
                [ -n "$a" ] && args+=("$a")
            done
            set -- "${args[@]}"
            break
        fi
    done
fi

# 检查容器是否存在
if [ "$num" -gt 0 ]; then
    echo "Container:[${CONTAINER_NAME}] already exists"

    # 检查是否有参数传入
    if [ "$#" -gt 0 ]; then
        param="$1"
        if [ "$param" = "d" ] || [ "$param" = "destroy" ]; then
            docker stop ${CONTAINER_NAME}
            docker rm -f ${CONTAINER_NAME}
            echo "container:[${CONTAINER_NAME}] finish destroy."
            exit
        fi

        if [ "$param" = "p" ] || [ "$param" = "print" ]; then
            echo "container:[${CONTAINER_NAME}]."
            exit
        fi
    fi
    
    # 如果是 -v 模式，显示信息后退出
    if [ "$show_connect_cmd" = true ]; then
        # 确保容器在运行状态
        container_status=$(docker inspect -f '{{.State.Status}}' ${CONTAINER_NAME} 2>/dev/null)
        if [ "$container_status" != "running" ]; then
            echo "Starting container..."
            docker start ${CONTAINER_NAME}
            check_command_result $?
            echo "Container:[${CONTAINER_NAME}] Started successfully!"
        fi
        show_container_info
        exit
    fi
    
    # 默认行为：进入容器
    docker start ${CONTAINER_NAME}        
    docker exec -ti --env COLUMNS=`tput cols` --env LINES=`tput lines` ${CONTAINER_NAME} /bin/bash
    exit
    
else
    echo "Container:[${CONTAINER_NAME}] does not exist, will be created and go inside it."

    eval ${command}
    check_command_result $?
    echo "Container:[${CONTAINER_NAME}] Created successfully!"

    docker start ${CONTAINER_NAME}
    exit_status=$?
    check_command_result $exit_status

    echo "Container:[${CONTAINER_NAME}] Started successfully!"
    
    # 如果是 -v 模式，显示信息后退出
    if [ "$show_connect_cmd" = true ]; then
        show_container_info
        exit
    fi
    
    docker exec -ti --env COLUMNS=`tput cols` --env LINES=`tput lines` ${CONTAINER_NAME} /bin/bash
    exit
fi