#!/bin/bash
#node_type主节点传入head,副节点传入worker
#model_name和obs桶中的权重文件名称一致
#ray_head_ip为主节点ip地址
export node_type=$1
export ray_head_ip=$2
export model_name=$3
export port=$4

export dir="/home/models-file/"
export image_id="swr.cn-north-9.myhuaweicloud.com/hw-sac/lm_inference:912"
export region=$(curl -s http://169.254.169.254/openstack/latest/meta_data.json | python3 -c "import sys, json; print(json.load(sys.stdin)['region_id'])")

# 检查Docker是否安装
echo "\n===== 检查Docker安装状态 ======"
if docker -v &>/dev/null; then
    echo "Docker已安装: $(docker -v | grep -o '[0-9.]*' | head -1)"
else
    echo "Docker未安装，开始安装..."
    # 安装Docker
    yum install -y docker-engine.aarch64 docker-engine-selinux.noarch docker-runc.aarch64
    if [ $? -eq 0 ]; then
        echo "Docker安装成功"
    else
        echo "Docker安装失败"
    fi
fi

# 合并磁盘
if ! mountpoint /home; then
    wget -O create_disk_partitions.sh https://documentation-samples.obs.cn-north-4.myhuaweicloud.com/solution-as-code-publicbucket/solution-as-code-moudle/one-click-deployment-ascend/userdata/create_disk_partitions.sh &&
        sh create_disk_partitions.sh &&
        rm create_disk_partitions.sh
fi

mkdir /home/models-file/
#安装docker-compose
if ! docker-compose --version &>/dev/null; then
    VERSION="v2.29.0"
    ARCH=$(uname -m)
    if [ "$ARCH" != "aarch64" ]; then
        echo "错误：此脚本仅适用于 ARM64 架构，当前架构为 $ARCH"
        exit 1
    fi

    echo "正在安装 Docker Compose $VERSION ..."
    wget -P /usr/local/bin/ https://documentation-samples.obs.cn-north-4.myhuaweicloud.com/solution-as-code-publicbucket/solution-as-code-moudle/one-click-deployment-ascend/userdata/docker-compose
    sudo chmod +x /usr/local/bin/docker-compose
    sudo ln -sf /usr/local/bin/docker-compose /usr/bin/docker-compose
    echo "✓ 安装完成: $(docker-compose --version)"
fi


# 同步权重
echo "============下载模型权重============"

# 创建模型特定的标识
model_download_flag="${dir}model-download-${model_name}.log"

if [ -f "${model_download_flag}" ]; then
    echo "✓ 模型 ${model_name} 已下载过，跳过"
else
    # 切换到目标目录，失败则退出
    cd "${dir}" || exit 1
    # 1. 准备下载工具
    echo "📦 准备下载工具..."
    rm -rf obs-sync
    if ! git clone --depth 1 https://gitee.com/flexus-agent/obs-sync.git; then
        echo "❌ 克隆仓库失败"
        exit 1
    fi

    cd obs-sync || exit 1
    cp .env.template .env.task

    sed -i \
        -e "s|^DEST_DIR=.*|DEST_DIR=${dir}|" \
        -e "s|^OBS_ENDPOINT=.*|OBS_ENDPOINT=obs.${region}.myhuaweicloud.com|" \
        -e 's|^CENTER_SERVER=.*|CENTER_SERVER=models-download.wangzhou3.top|' \
        -e 's|^CENTER_SERVER_HTTP_SCHEMA=.*|CENTER_SERVER_HTTP_SCHEMA=http|' \
        .env.task
    echo "✓ 配置文件已修改: .env.task"
    grep -E "^(DEST_DIR|OBS_ENDPOINT|CENTER_SERVER|CENTER_SERVER_HTTP_SCHEMA)=" .env.task

    cp docker-compose.yml.template docker-compose.yml
    sed -i 's/sac-models-download:0\.2/sac-models-download:0.3/g' ${dir}/obs-sync/docker-compose.yml

    # 2. 启动容器
    if ! docker-compose --env-file .env.task up -d obs-sync-task; then
        echo "❌ 启动容器失败"
        docker-compose --env-file .env.task down obs-sync-task 2>/dev/null || true
        exit 1
    fi

    echo "⏳ 等待容器启动..."
    sleep 5

    # 3. 获取任务ID
    echo "🔍 获取下载任务ID..."
    TASK_ID=""
    api_retry=0
    max_api_retries=5

    while [ -z "$TASK_ID" ] && [ $api_retry -lt $max_api_retries ]; do
        api_retry=$((api_retry + 1))

        response=$(curl -s --max-time 10 -X POST "http://127.0.0.1/api/task?modelName=${model_name}&taskType=local" 2>&1)
        curl_exit_code=$?

        if [ $curl_exit_code -ne 0 ] || [ -z "$response" ]; then
            echo "⚠ API 请求失败 (尝试 $api_retry/$max_api_retries)"
            echo "退出码: $curl_exit_code"
            echo "错误详情: $response"
            sleep 2
            continue
        fi

        # 检查 code 是否为 200
        if echo "$response" | grep -q '"code":200'; then
            TASK_ID=$(echo "$response" | sed -n 's/.*"data":"\([^"]*\)".*/\1/p')
            if [ -n "$TASK_ID" ] && [ "$TASK_ID" != "null" ]; then
                echo "✓ 任务ID: $TASK_ID"
                break
            else
                # code=200 但 data 为空或 null
                echo "⚠ 获取任务ID失败 (尝试 $api_retry/$max_api_retries)"
                echo "响应中 data 字段为空或 null"
                echo "完整响应: $response"
                [ $api_retry -lt $max_api_retries ] && sleep 2
            fi
        else
            # 提取并打印 msg
            msg=$(echo "$response" | sed -n 's/.*"msg":"\([^"]*\)".*/\1/p')
            echo "⚠ 获取任务ID失败 (尝试 $api_retry/$max_api_retries)"
            echo "错误信息: $msg"
            [ $api_retry -lt $max_api_retries ] && sleep 2
        fi
    done

    if [ -z "$TASK_ID" ]; then
        echo "❌ 无法获取任务ID，已重试 $max_api_retries 次"
        docker logs --tail 20 obs-sync-task 2>&1 || true
        docker-compose --env-file .env.task down obs-sync-task
        exit 1
    fi

    # 4. 监控任务状态
    echo "📊 开始监控下载进度..."
    retry_count=0
    max_retry_count=5
    monitor_interval=5
    max_wait_seconds=7200  # 2小时超时
    waited_seconds=0

    while [ $waited_seconds -lt $max_wait_seconds ]; do
        # 获取任务状态，添加超时保护
        status_response=$(curl -s --max-time 5 "http://127.0.0.1/api/task?id=$TASK_ID" 2>&1)

        if [ $? -ne 0 ] || [ -z "$status_response" ]; then
            echo "$(date '+%H:%M:%S') ⚠ 无法获取任务状态，API 可能异常，继续等待..."
            sleep $monitor_interval
            waited_seconds=$((waited_seconds + monitor_interval))
            continue
        fi

        # 提取状态和进度
        STATUS=$(echo "$status_response" | sed -n 's/.*"status"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')
        SizeProgress=$(echo "$status_response" | sed -n 's/.*"sizeProgress"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')

        # 备用提取方法
        if [ -z "$STATUS" ]; then
            STATUS=$(echo "$status_response" | grep -o '"status":"[^"]*"' | head -1 | cut -d'"' -f4)
        fi
        if [ -z "$SizeProgress" ]; then
            SizeProgress=$(echo "$status_response" | grep -o '"sizeProgress":"[^"]*"' | head -1 | cut -d'"' -f4)
        fi

        # 默认值
        [ -z "$SizeProgress" ] && SizeProgress="未知"

        echo "$(date '+%H:%M:%S') 状态: ${STATUS:-未知} - 进度: ${SizeProgress}"

        # 检查是否完成
        if [ "$STATUS" = "FINISHED" ]; then
            echo "✅ 下载完成！"
            break
        elif [ "$STATUS" = "FAILED" ]; then
            if [ $retry_count -ge $max_retry_count ]; then
                echo "❌ 下载失败，已达最大重试次数 ($max_retry_count)"
                docker logs --tail 30 obs-sync-task 2>&1 || true
                docker-compose --env-file .env.task down obs-sync-task
                exit 1
            fi
            retry_count=$((retry_count + 1))
            echo "⚠ 下载失败，正在重试 (第 $retry_count/$max_retry_count 次)"
            curl -s -X POST "http://127.0.0.1/api/task/retry?id=${TASK_ID}" > /dev/null
            sleep 10  # 重试后等待更长时间
            continue
        elif [ -z "$STATUS" ]; then
            echo "$(date '+%H:%M:%S') ⚠ 状态异常，继续等待..."
            sleep $monitor_interval
        else
            # PENDING、RUNNING 等其他状态，正常等待
            sleep $monitor_interval
        fi

        waited_seconds=$((waited_seconds + monitor_interval))
    done

    # 检查是否超时
    if [ $waited_seconds -ge $max_wait_seconds ]; then
        echo "❌ 下载超时（超过 $((max_wait_seconds/3600)) 小时）"
        docker-compose --env-file .env.task down obs-sync-task
        exit 1
    fi

    # 5. 清理并标记完成
    echo "🧹 清理容器..."
    docker-compose --env-file .env.task down obs-sync-task 2>/dev/null || true
    docker stop obs-sync-task 2>/dev/null || true
    docker rm obs-sync-task 2>/dev/null || true
    echo "✅ 容器已清理"
    echo "finished" > "${model_download_flag}"
    echo "✅ 模型 ${model_name} 下载完成"
fi
# 获取NPU驱动版本
echo "===== 检查NPU驱动版本 ======"
version=$(npu-smi info -t board -i 1 2>/dev/null | grep -i "software\|firmware" | grep -oE "[0-9]+\.[0-9]+\.[0-9]+" | head -1)
# 检查是否获取到版本号
if [ -z "$version" ]; then
    echo "错误：无法获取NPU驱动版本"
else
    echo "当前驱动版本: $version"

    # 版本比较
    min_version="24.1.0"
    # 使用sort进行版本比较
    if printf '%s\n%s\n' "$version" "$min_version" | sort -V | head -n1 | grep -q "^$min_version$"; then
        echo "✅ 版本符合要求（≥ $min_version）"
    else
        echo "❌ 版本过低，需要更新到 $min_version 或更高"
    fi
fi

# 检查卡是否被占用
command -v npu-smi &>/dev/null || {
    echo "错误: npu-smi 命令未找到"
    exit 1
}

# 获取被占用的NPU卡
npu_info=$(npu-smi info 2>/dev/null)
occupied_cards=""

# 解析每一行，查找有进程ID的卡
while IFS= read -r line; do
    # 匹配格式: | 0       0                 | 199676        | python
    if [[ $line =~ ^\|\ ([0-9]+)[[:space:]]+[0-9]+[[:space:]]*\|[[:space:]]*[0-9]+ ]]; then
        card_num=$(echo "$line" | awk -F'|' '{print $2}' | awk '{print $1}')
        occupied_cards="$occupied_cards $card_num"
    fi
done <<<"$npu_info"

# 去重并排序被占用的卡
occupied_cards=$(echo $occupied_cards | tr ' ' '\n' | sort -n | uniq | xargs)

# 找出未被占用的卡 (0-7)
free_cards=""
for i in {0..7}; do
    if [[ ! " $occupied_cards " =~ " $i " ]]; then
        free_cards="$free_cards $i"
    fi
done
free_cards=$(echo $free_cards | xargs)

# 统计数量
occupied_count=$(echo $occupied_cards | wc -w)
free_count=$(echo $free_cards | wc -w)

# 输出结果
if [ -n "$occupied_cards" ]; then
    echo "被占用的NPU卡 ($occupied_count 张): $occupied_cards"
    echo "未被占用的NPU卡 ($free_count 张): $free_cards"
else
    echo "8张NPU卡均未被占用: 0 1 2 3 4 5 6 7"
fi

# 部署模型
# 定义文件路径
RUN_VLLM_SCRIPT="${dir}run_vllm.sh"
RAY_WORKER_CHECKER="${dir}ray_worker_checker.py"

# 下载 run_vllm.sh（如果不存在或需要更新）
if [ -f "$RUN_VLLM_SCRIPT" ]; then
    echo "✓ run_vllm.sh 已存在，跳过下载"
else
    echo "下载 run_vllm.sh..."
    wget -O ${RUN_VLLM_SCRIPT} https://documentation-samples-17.obs.cn-north-9.myhuaweicloud.com/solution-as-code-publicbucket/solution-as-code-module/quickly-deploy-llm-on-modelarts-lite-devserver/userdata/deploy-vl-model/dual-machine/run_vllm.sh
    if [ $? -eq 0 ]; then
        echo "✓ run_vllm.sh 下载成功"
    else
        echo "✗ run_vllm.sh 下载失败"
        exit 1
    fi
fi

# 下载 ray_worker_checker.py（如果不存在或需要更新）
if [ -f "$RAY_WORKER_CHECKER" ]; then
    echo "✓ ray_worker_checker.py 已存在，跳过下载"
else
    echo "下载 ray_worker_checker.py..."
    wget -O ${RAY_WORKER_CHECKER} https://documentation-samples-17.obs.cn-north-9.myhuaweicloud.com/solution-as-code-publicbucket/solution-as-code-module/quickly-deploy-llm-on-modelarts-lite-devserver/userdata/deploy-vl-model/dual-machine/ray_worker_checker.py
    if [ $? -eq 0 ]; then
        echo "✓ ray_worker_checker.py 下载成功"
    else
        echo "✗ ray_worker_checker.py 下载失败"
        exit 1
    fi
fi

# 设置执行权限
chmod 755 ${RUN_VLLM_SCRIPT}
chmod 755 ${RAY_WORKER_CHECKER}
docker pull ${image_id}

echo "运行Docker容器..."
container_name="${model_name}_$(date +%Y%m%d_%H%M%S)_${RANDOM}"

# 启动容器
docker run -itd \
    --device=/dev/davinci0 \
    --device=/dev/davinci1 \
    --device=/dev/davinci2 \
    --device=/dev/davinci3 \
    --device=/dev/davinci4 \
    --device=/dev/davinci5 \
    --device=/dev/davinci6 \
    --device=/dev/davinci7 \
    --privileged \
    -v /etc/localtime:/etc/localtime \
    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
    -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \
    --device=/dev/davinci_manager \
    --device=/dev/devmm_svm \
    --device=/dev/hisi_hdc \
    -v /var/log/npu/:/usr/slog \
    -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \
    -v /sys/fs/cgroup:/sys/fs/cgroup:ro \
    -v ${dir}:${dir} \
    --net=host \
    --name ${container_name} \
    ${image_id} \
    /bin/bash

# 检查容器是否启动成功
if [ $? -ne 0 ]; then
    echo "容器启动失败"
    exit 1
fi

# 等待服务就绪
echo "等待服务就绪..."
MAX_RETRIES=60
for i in $(seq 1 $MAX_RETRIES); do
    # 方法1：检查容器日志
    if docker logs ${container_name} 2>&1 | grep -q "Application startup complete\|Uvicorn running"; then
        echo "服务启动成功！"
        break
    fi
done
docker exec -d -u root ${container_name} bash -l -c "
    cd ${dir}
    nohup bash run_vllm.sh ${node_type} ${ray_head_ip} ${model_name} ${port} > vllm.log 2>&1 &
    chmod 755 vllm.log
    echo \$! > vllm.pid
"

if [ "$node_type" = "head" ]; then
    echo "node_type 是 head，开始检查 ${port} 端口..."

    # 监听端口的代码
    echo "等待${port}端口启动..."
    attempt=1
    max_attempts=360 # 30分钟 / 5秒 = 360次

    while true; do
        echo "------ 动态日志（第 $attempt 次检查）------"
        if [ -f "${dir}/vllm.log" ]; then
            tail -n 20 "${dir}/vllm.log" 2>/dev/null || echo "无法读取日志文件"
        else
            # 如果宿主机没有日志文件，回退到 docker logs
            docker logs --tail 10 "${container_name}" 2>&1 || true
        fi
        if docker exec ${container_name} netstat -tlnp 2>/dev/null | grep -q ":${port} "; then
            echo ""
            echo "✅ 模型部署成功！"
            echo "端口${port}已监听"
            echo "服务信息:"
            docker exec ${container_name} netstat -tlnp | grep ${port}
            break
        fi

        echo -n "."
        echo "tail -f ${dir}vllm.log"
        sleep 5

        attempt=$((attempt + 1))
        if [ $attempt -gt $max_attempts ]; then
            echo ""
            echo "❌ 部署超时（30分钟），端口8080未能正常启动"
            exit 1
        fi
    done
else
    echo "node_type 不是 head，退出脚本。请切换至主节点查看日志。"
    exit 0 # 或 exit 1，根据需要决定是否正常退出
fi