#!/bin/bash
# manage_cann_pkg.sh - SOC Type Supported Version

# === 配置 ===
CANN_VERSION_INFO="/home/developer/Ascend/ascend-toolkit/latest/share/info/acl_extend/version.info"
CURRENT_VERSION_CACHE="$HOME/.cann/current_version.json"
VERSIONS_URL="https://mirror-centralrepo.devcloud.cn-north-4.huaweicloud.com/artifactory/cann-run-mirror/versions.json"
CANN_TASK_FILE="$HOME/.cann/upgrade_task.json"
TEMP_DIR="/tmp/cann_upgrade_$$"

# 自动检测架构
ARCH=$(uname -m)
if [[ "$ARCH" == "aarch64" ]]; then
    TARGET_ARCH="aarch64"
elif [[ "$ARCH" == "x86_64" ]]; then
    TARGET_ARCH="x86_64"
else
    echo "{\"error_code\":\"128\",\"error_msg\":\"Unsupported architecture: $ARCH\"}"
    exit 128
fi

# 默认SOC类型
SOC_TYPE="A2"
# 映射表：SOC类型 -> Ops包字段名
declare -A OPS_MAP=( ["A2"]="Ascend-cann-910b_ops_url" ["A3"]="Ascend-cann-A3_ops_url" ["A5"]="Ascend-cann-950_ops_url" )

mkdir -p "$(dirname "$CANN_TASK_FILE")" "$TEMP_DIR" 2>/dev/null || true
trap 'rm -rf "$TEMP_DIR"' EXIT

# === 工具函数 ===
log_error() { echo "$1" >&2; }
output_json() { echo "$1"; }
fail() { output_json "{\"error_code\":\"$1\",\"error_msg\":\"$2\"}"; exit "$1"; }

# 检查依赖
for cmd in jq wget uname date ps pkill sudo; do
    if ! command -v "$cmd" &> /dev/null; then
        fail 127 "Required command '$cmd' is not installed."
    fi
done

# === 参数解析 (新增 --soc 处理) ===
ACTION=""
TARGET_VERSION=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --version)
            ACTION="version"
            # 检查是否有指定 SOC 参数（可选）
            shift
            ;;
        --list)
            ACTION="list"
            shift
            ;;
        --upgrade)
            ACTION="upgrade"
            TARGET_VERSION="$2"
            [[ -z "$TARGET_VERSION" ]] && fail 2 "Missing version for --upgrade"
            shift 2
            ;;
        --forceupgrade)
            ACTION="forceupgrade"
            TARGET_VERSION="$2"
            [[ -z "$TARGET_VERSION" ]] && fail 3 "Missing version for --forceupgrade"
            shift 2
            ;;
        --soc)
            SOC_PARAM="$2"
            if [[ "$SOC_PARAM" =~ ^(A2|A3|A5)$ ]]; then
                SOC_TYPE="$SOC_PARAM"
            else
                fail 4 "Invalid SOC type: $SOC_PARAM. Use A2, A3, or A5."
            fi
            shift 2
            ;;
        *)
            fail 1 "Unknown option: $1"
            ;;
    esac
done

[[ -z "$ACTION" ]] && fail 2 "No action specified"

# 获取当前 SOC 对应的 Ops 字段名
OPS_FIELD_NAME="${OPS_MAP[$SOC_TYPE]}"
if [[ -z "$OPS_FIELD_NAME" ]]; then
    fail 4 "Unsupported SOC type configuration: $SOC_TYPE"
fi

# === 强制清理函数 ===
force_cleanup_previous_upgrade() {
    local ASCEND_ROOT="/home/developer/Ascend"
    local ASCEND_BK="${ASCEND_ROOT}.bk"

    local pids_to_kill=""
    while IFS= read -r pid; do
        if [[ -f "/proc/$pid/cmdline" ]]; then
            cmdline=$(tr '\0' ' ' < "/proc/$pid/cmdline")
            if [[ "$cmdline" =~ /tmp/cann_upgrade_async_[0-9]+/ ]]; then
                pids_to_kill="$pids_to_kill $pid"
            fi
        fi
    done < <(pgrep -f '/tmp/cann_upgrade_async_' 2>/dev/null || true)

    for pid in $pids_to_kill; do
        kill "$pid" 2>/dev/null || true
    done
    sleep 0.3
    for pid in $pids_to_kill; do
        if kill -0 "$pid" 2>/dev/null; then
            kill -9 "$pid" 2>/dev/null || true
        fi
    done

    if [[ -e "$ASCEND_BK" ]]; then
        [[ -e "$ASCEND_ROOT" ]] && sudo rm -rf "$ASCEND_ROOT"
        if ! sudo mv "$ASCEND_BK" "$ASCEND_ROOT" 2>/dev/null; then
            sudo rm -rf "$ASCEND_BK"
        fi
    fi
    sudo rm -f "$CANN_TASK_FILE"
    sudo rm -rf /tmp/cann_upgrade_* /tmp/cann_upgrade_async_* 2>/dev/null || true
    [[ -t 1 ]] && stty sane 2>/dev/null || true
}

# === 执行动作 ===
case "$ACTION" in
    "list")
        # 下载完整版本定义，不进行过滤
        wget -q -O "$TEMP_DIR/versions.json" "$VERSIONS_URL" --no-check-certificate || fail 30 "Failed to download versions.json"
        
        # 直接输出完整的 versions.json 内容
        cat "$TEMP_DIR/versions.json"
        ;;
    "version")
        # 下载版本定义
        wget -q -O "$TEMP_DIR/versions.json" "$VERSIONS_URL" --no-check-certificate || fail 30 "Failed to download versions.json"

        # --- 核心修改：根据 SOC_TYPE 过滤版本列表 ---
        # 过滤逻辑：仅保留 download_url 中存在对应 SOC 的 Ops 包字段的版本
        FILTERED_VERSIONS_FILE="$TEMP_DIR/filtered_versions.json"

        # 使用 jq 进行过滤
        # 使用 jq 进行过滤，保留非 versions 字段，只过滤 versions 数组
        jq -c --arg arch "$TARGET_ARCH" --arg ops_field "$OPS_FIELD_NAME" '
            {
                master_branch: .master_branch,
                special_branch: .special_branch,
                versions: [
                    .versions[] |
                    select(
                        .download_url[$arch] != null and
                        .download_url[$arch][$ops_field] != null and
                        .download_url[$arch][$ops_field] != "null" and
                        .download_url[$arch][$ops_field] != ""
                    )
                ]
            }
         ' "$TEMP_DIR/versions.json" > "$FILTERED_VERSIONS_FILE"

        # 如果过滤后为空，创建一个空数组防止报错
        if ! jq -e '.versions' "$FILTERED_VERSIONS_FILE" >/dev/null; then
            echo '{"versions":[]}' > "$FILTERED_VERSIONS_FILE"
        fi

        VERSIONS_FILE="$FILTERED_VERSIONS_FILE"
        LOCAL_VERSION=""
        TIMESTAMP_RAW=""

        if [[ -f "$CURRENT_VERSION_CACHE" ]]; then
            LOCAL_VERSION=$(jq -r '.Version // empty' "$CURRENT_VERSION_CACHE")
            TIMESTAMP_RAW=$(jq -r '.timestamp // empty' "$CURRENT_VERSION_CACHE")
        fi
        if [[ -z "$LOCAL_VERSION" || -z "$TIMESTAMP_RAW" ]] && [[ -f "$CANN_VERSION_INFO" ]]; then
            LOCAL_VERSION=$(grep "^Version=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
            TIMESTAMP_RAW=$(grep "^timestamp=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
        fi

        if [[ -z "$LOCAL_VERSION" || -z "$TIMESTAMP_RAW" ]]; then
            LOCAL_VERSION="unknown"
            TIMESTAMP_RAW="19700101000000"
        fi

        # --- 时间格式化与显示逻辑 (保持原逻辑不变) ---
        CLEAN_TS=${TIMESTAMP_RAW/_/}
        if [[ ${#CLEAN_TS} -ge 8 ]]; then LOCAL_DATE=${CLEAN_TS:0:8}; else LOCAL_DATE="19700101"; fi
        if [[ "$LOCAL_VERSION" != "unknown" ]] && [[ ${#LOCAL_DATE} -ne 8 ]]; then
            fail 12 "Invalid timestamp format"
        fi

        #cat $VERSIONS_FILE
        MASTER_BRANCH=$(jq -r '.master_branch // empty' "$VERSIONS_FILE")
        if [[ -z "$MASTER_BRANCH" ]]; then fail 35 "Field 'master_branch' missing in versions.json"; fi

        # === 关键逻辑：判断是否属于 special_branch 或 构建显示信息 ===
        # (此处保留原脚本中从 SPECIAL_BRANCHES 定义到 CURRENT_JSON 生成的所有代码)
        # 为了防止遗漏，我将原脚本的这部分逻辑完整保留在此处：
        SPECIAL_BRANCHES=$(jq -r '.special_branch // [] | .[]' "$VERSIONS_FILE")
        IS_INTERMEDIATE=false
        MATCHED_INDEX=-1
        if [[ "$LOCAL_VERSION" == "unknown" ]]; then
            CURRENT_DISPLAY="not installed"
            TIME_ISO="1970-01-01T00:00:00Z"
            CURRENT_DESC="CANN Toolkit not installed"
        elif [[ "$LOCAL_VERSION" == "$MASTER_BRANCH" ]]; then
            # === 关键逻辑：判断是否属于 special_branch ===
            CANDIDATE_VERSION="${LOCAL_VERSION}-${LOCAL_DATE}"
            IS_SPECIAL=false
            for sb in $SPECIAL_BRANCHES; do
                if [[ "$sb" == "$CANDIDATE_VERSION" ]]; then
                    IS_SPECIAL=true
                    break
                fi
            done
            if [[ "$IS_SPECIAL" == "true" ]]; then
                # 是 special_branch 中的特殊版本
                CURRENT_DISPLAY="$CANDIDATE_VERSION"
                CLEAN_FULL=${TIMESTAMP_RAW/_/}
                if [[ ${#CLEAN_FULL} -ge 14 ]]; then
                    YEAR=${CLEAN_FULL:0:4}; MONTH=${CLEAN_FULL:4:2}; DAY=${CLEAN_FULL:6:2}
                    HOUR=${CLEAN_FULL:8:2}; MIN=${CLEAN_FULL:10:2}; SEC=${CLEAN_FULL:12:2}
                else
                    YEAR=1970; MONTH=01; DAY=01; HOUR=00; MIN=00; SEC=00
                fi
                TIME_ISO="${YEAR}-${MONTH}-${DAY}T${HOUR}:${MIN}:${SEC}Z"
                CURRENT_DESC="特殊发布版本 $CANDIDATE_VERSION"
                # 尝试从 versions 列表获取 desc 和 index
                MATCHED_ENTRY=$(jq --arg ver "$CANDIDATE_VERSION" ' .versions | map(select(.version == $ver)) | if length > 0 then .[0] else null end ' "$VERSIONS_FILE")
                if [[ "$MATCHED_ENTRY" != "null" ]] && [[ -n "$MATCHED_ENTRY" ]]; then
                    CURRENT_DESC=$(echo "$MATCHED_ENTRY" | jq -r '.desc')
                    MATCHED_INDEX=$(jq --arg ver "$CANDIDATE_VERSION" ' .versions | to_entries | map(select(.value.version == $ver)) | if length > 0 then .[0].key | tonumber else -1 end ' "$VERSIONS_FILE")
                else
                    MATCHED_INDEX=-1
                fi
            else
                # 不是 special_branch，走标准 master 逻辑
                CURRENT_DISPLAY="master-$LOCAL_DATE"
                CLEAN_FULL=${TIMESTAMP_RAW/_/}
                if [[ ${#CLEAN_FULL} -ge 14 ]]; then
                    YEAR=${CLEAN_FULL:0:4}; MONTH=${CLEAN_FULL:4:2}; DAY=${CLEAN_FULL:6:2}
                    HOUR=${CLEAN_FULL:8:2}; MIN=${CLEAN_FULL:10:2}; SEC=${CLEAN_FULL:12:2}
                else
                    YEAR=1970; MONTH=01; DAY=01; HOUR=00; MIN=00; SEC=00
                fi
                TIME_ISO="${YEAR}-${MONTH}-${DAY}T${HOUR}:${MIN}:${SEC}Z"
                CURRENT_DESC="主线每周构建版本"
		MATCHED_INDEX=0
            fi
        else
            # 非 master 分支：查找 exact version
            MATCHED_ENTRY=$(jq --arg ver "$LOCAL_VERSION" ' .versions | map(select(.version == $ver)) | if length > 0 then .[0] else null end ' "$VERSIONS_FILE")
            if [[ "$MATCHED_ENTRY" == "null" ]] || [[ -z "$MATCHED_ENTRY" ]]; then
                # versions.json 中无此 version
                CURRENT_DISPLAY="${LOCAL_VERSION}-${LOCAL_DATE}"
                CLEAN_FULL=${TIMESTAMP_RAW/_/}
                if [[ ${#CLEAN_FULL} -ge 14 ]]; then
                    YEAR=${CLEAN_FULL:0:4}; MONTH=${CLEAN_FULL:4:2}; DAY=${CLEAN_FULL:6:2}
                    HOUR=${CLEAN_FULL:8:2}; MIN=${CLEAN_FULL:10:2}; SEC=${CLEAN_FULL:12:2}
                else
                    YEAR=1970; MONTH=01; DAY=01; HOUR=00; MIN=00; SEC=00
                fi
                TIME_ISO="${YEAR}-${MONTH}-${DAY}T${HOUR}:${MIN}:${SEC}Z"
                CURRENT_DESC="$LOCAL_VERSION 中间迭代版本"
                IS_INTERMEDIATE=true
            else
                REMOTE_TIME_ISO=$(echo "$MATCHED_ENTRY" | jq -r '.time // empty')
                # 从 remote time 提取日期 YYYYMMDD
                if [[ "$REMOTE_TIME_ISO" =~ ^([0-9]{4})-([0-9]{2})-([0-9]{2})T ]]; then
                    REMOTE_DATE="${BASH_REMATCH[1]}${BASH_REMATCH[2]}${BASH_REMATCH[3]}"
                else
                    REMOTE_DATE=""
                fi
                if [[ "$REMOTE_DATE" == "$LOCAL_DATE" ]]; then
                    # 日期匹配 → 正式版本
                    CURRENT_DISPLAY="$LOCAL_VERSION"
                    CLEAN_FULL=${TIMESTAMP_RAW/_/}
                    if [[ ${#CLEAN_FULL} -ge 14 ]]; then
                        YEAR=${CLEAN_FULL:0:4}; MONTH=${CLEAN_FULL:4:2}; DAY=${CLEAN_FULL:6:2}
                        HOUR=${CLEAN_FULL:8:2}; MIN=${CLEAN_FULL:10:2}; SEC=${CLEAN_FULL:12:2}
                    else
                        YEAR=1970; MONTH=01; DAY=01; HOUR=00; MIN=00; SEC=00
                    fi
                    TIME_ISO="${YEAR}-${MONTH}-${DAY}T${HOUR}:${MIN}:${SEC}Z"
                    CURRENT_DESC=$(echo "$MATCHED_ENTRY" | jq -r '.desc')
                    MATCHED_INDEX=$(jq --arg ver "$LOCAL_VERSION" ' .versions | to_entries | map(select(.value.version == $ver)) | if length > 0 then .[0].key | tonumber else -1 end ' "$VERSIONS_FILE")
                else
                    # 日期不匹配 → 中间迭代版本
                    CURRENT_DISPLAY="${LOCAL_VERSION}-${LOCAL_DATE}"
                    CLEAN_FULL=${TIMESTAMP_RAW/_/}
                    if [[ ${#CLEAN_FULL} -ge 14 ]]; then
                        YEAR=${CLEAN_FULL:0:4}; MONTH=${CLEAN_FULL:4:2}; DAY=${CLEAN_FULL:6:2}
                        HOUR=${CLEAN_FULL:8:2}; MIN=${CLEAN_FULL:10:2}; SEC=${CLEAN_FULL:12:2}
                    else
                        YEAR=1970; MONTH=01; DAY=01; HOUR=00; MIN=00; SEC=00
                    fi
                    TIME_ISO="${YEAR}-${MONTH}-${DAY}T${HOUR}:${MIN}:${SEC}Z"
                    CURRENT_DESC="$LOCAL_VERSION 中间迭代版本"
                    IS_INTERMEDIATE=true
                    MATCHED_INDEX=$(jq --arg ver "$LOCAL_VERSION" ' .versions | to_entries | map(select(.value.version == $ver)) | if length > 0 then .[0].key | tonumber else -1 end ' "$VERSIONS_FILE")
                fi
            fi
        fi

        # 构建当前版本 JSON
        CURRENT_JSON=$(jq -n \
            --arg ver "$CURRENT_DISPLAY" \
            --arg time "$TIME_ISO" \
            --arg desc "$CURRENT_DESC" \
            '{version: $ver, time: $time, desc: $desc}')

        # 3. 基于新规则构建升级列表 (核心修改点)
        # 规则：如果本地版本是 unknown，或者本地版本不在 version.json 中 (MATCHED_INDEX == -1)，则返回所有版本
        if [[ "$LOCAL_VERSION" == "unknown" ]] || [[ "$MATCHED_INDEX" -eq -1 ]]; then
            UPGRADE_LIST_JSON=$(jq -c '[.versions[] | {version: .version, time: .time, desc: .desc}]' "$VERSIONS_FILE")
        else
            # 规则：本地版本存在且在列表中 -> 保持原有逻辑 (仅显示新版本+3个旧版本)
            # 如果是中间版本 (IS_INTERMEDIATE)，原脚本逻辑通常返回空或特定处理
            # 根据你的原脚本逻辑：如果是中间版本且 MATCHED_INDEX >=0
            if [[ "$IS_INTERMEDIATE" == "true" ]]; then
                # 原脚本逻辑：如果是中间版本，通常只返回比它旧的版本或者空
                # 这里我们保持原样：返回空列表 (或者你可以选择返回所有，根据需求调整)
                # 为了严格符合"其他逻辑不变"，我们复现原脚本的中间版本逻辑
                ALL_VERSIONS=$(jq ' .versions | map({ version: .version, time: .time, desc: .desc }) ' "$VERSIONS_FILE")
                IDX=$MATCHED_INDEX
                # 原脚本逻辑：newer_or_self (包含自己及之前) + older_three (之后的3个)
                NEWER_OR_SELF=$(echo "$ALL_VERSIONS" | jq --argjson i "$IDX" '.[0:$i+1]')
                OLDER_THREE=$(echo "$ALL_VERSIONS" | jq --argjson i "$IDX" '.[$i+1:] | .[0:3]')
                UPGRADE_LIST_JSON=$(jq -n --argjson newer "$NEWER_OR_SELF" --argjson older "$OLDER_THREE" '$newer + $older')
            else
                # 标准正式版本逻辑：返回比它新的所有版本 + 比它旧的3个版本
                INDEX=$(jq --arg ver "$CURRENT_DISPLAY" ' .versions | to_entries | map(select(.value.version == $ver)) | if length > 0 then .[0].key | tonumber else -1 end ' "$VERSIONS_FILE")
                if [[ "$INDEX" -eq -1 ]]; then
                    UPGRADE_LIST_JSON=$(jq -n '[]')
                else
                    ALL_VERSIONS=$(jq ' .versions | map({ version: .version, time: .time, desc: .desc }) ' "$VERSIONS_FILE")
                    IDX=$INDEX
                    NEWER=$(echo "$ALL_VERSIONS" | jq --argjson i "$IDX" '.[0:$i]')
                    OLDER=$(echo "$ALL_VERSIONS" | jq --argjson i "$IDX" '.[$i+1:] | .[0:3]')
                    UPGRADE_LIST_JSON=$(jq -n --argjson newer "$NEWER" --argjson older "$OLDER" '$newer + $older')
                fi
            fi
        fi

        # 4. 检查是否有进行中的任务 (保持原逻辑不变)
        TASK_JSON="null"
        if [[ -f "$CANN_TASK_FILE" ]]; then
            TASK_CONTENT=$(cat "$CANN_TASK_FILE")
            STATUS=$(echo "$TASK_CONTENT" | jq -r '.status // "1"')
            if [[ "$STATUS" == "1" ]]; then
                TASK_JSON=$(echo "$TASK_CONTENT" | jq '.process = "100%"')
            elif [[ "$STATUS" == "0" ]]; then
                START_TIME_STR=$(echo "$TASK_CONTENT" | jq -r '.start_time // empty')
                if [[ -n "$START_TIME_STR" && "$START_TIME_STR" != "null" ]]; then
                    Y=${START_TIME_STR:0:4}
                    M=${START_TIME_STR:4:2}
                    D=${START_TIME_STR:6:2}
                    h=${START_TIME_STR:9:2}
                    m=${START_TIME_STR:12:2}
                    s=${START_TIME_STR:15:2}
                    ISO_TIME="${Y}-${M}-${D}T${h}:${m}:${s}Z"
                    START_TS=$(date -d "$ISO_TIME" +%s 2>/dev/null || echo 0)
                    NOW_TS=$(date -u +%s)
                    ELAPSED=$((NOW_TS - START_TS))
                    if [[ $ELAPSED -lt 0 ]]; then
                        PROGRESS="0%"
                    elif [[ $ELAPSED -ge 900 ]]; then
                        PROGRESS="99%"
                    else
                        PERCENT=$((ELAPSED * 100 / 900))
                        if [[ $PERCENT -ge 99 ]]; then
                            PROGRESS="99%"
                        else
                            PROGRESS="${PERCENT}%"
                        fi
                    fi
                    TASK_JSON=$(echo "$TASK_CONTENT" | jq --arg p "$PROGRESS" '.process = $p')
                else
                    TASK_JSON="$TASK_CONTENT"
                fi
            else
                TASK_JSON="$TASK_CONTENT"
            fi
        fi

        # 5. 组装最终响应
        RESPONSE=$(jq -n \
            --argjson current "$CURRENT_JSON" \
            --argjson list "$UPGRADE_LIST_JSON" \
            --argjson task "$TASK_JSON" \
            '{ error_code: "0", error_msg: "success", current_version_info: $current, version_list: $list, task: $task }')
        output_json "$RESPONSE"
        ;;
    "upgrade")
        if [[ -f "$CANN_TASK_FILE" ]]; then
            STATUS=$(jq -r '.status // empty' "$CANN_TASK_FILE" 2>/dev/null || echo "")
            [[ "$STATUS" == "0" ]] && fail 21 "Upgrade task is running. Use --forceupgrade to override."
        fi

        output_json '{"error_code":"0","error_msg":"Upgrade task started"}'

        (
            # === 关键：临时目录由当前用户创建，不用 sudo ===
            TEMP_DIR_ASYNC="/tmp/cann_upgrade_async_$$"
            mkdir -p "$TEMP_DIR_ASYNC" 2>/dev/null || true
            trap 'rm -rf "$TEMP_DIR_ASYNC"' EXIT

            # --- 先下载 versions.json，保留错误可见性 ---
            if ! wget -q -O "$TEMP_DIR_ASYNC/versions.json" "$VERSIONS_URL" \
                    --no-check-certificate --timeout=20 --tries=2; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "30 download versions.json failed"
}
EOF
                exit 30
            fi

            # --- 从现在开始静默输出 ---
            exec >/dev/null 2>&1

            TARGET_ENTRY=$(jq -e --arg v "$TARGET_VERSION" '.versions[] | select(.version == $v)' "$TEMP_DIR_ASYNC/versions.json" 2>/dev/null) || true
            if [[ -z "$TARGET_ENTRY" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "20 target version not found"
}
EOF
                exit 20
            fi

            TOOLKIT_URL=$(echo "$TARGET_ENTRY" | jq -r ".download_url.${TARGET_ARCH}[\"Ascend-cann-toolkit_url\"]")
            # --- 核心修改：根据 SOC_TYPE 选择 Ops URL ---
            OPS_URL=$(echo "$TARGET_ENTRY" | jq -r ".download_url.${TARGET_ARCH}[\"$OPS_FIELD_NAME\"]")

            if [[ -z "$TOOLKIT_URL" || "$TOOLKIT_URL" == "null" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "22 toolkit URL missing for $TARGET_ARCH"
}
EOF
                exit 22
            fi
            if [[ -z "$OPS_URL" || "$OPS_URL" == "null" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "23 ops URL missing for $TARGET_ARCH"
}
EOF
                exit 23
            fi

            START_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
            cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 0,
    "target_version": "$TARGET_VERSION",
    "start_time": "$START_TIME",
    "end_time": "",
    "process": "0%",
    "remark": ""
}
EOF

            RUN_DIR="$TEMP_DIR_ASYNC/runs"
            mkdir -p "$RUN_DIR" 2>/dev/null || true
            if ! wget -q -O "$RUN_DIR/toolkit.run" "$TOOLKIT_URL" --no-check-certificate --timeout=60 --tries=2; then
                jq --arg remark "40 download run pkg failed" \
                   '.status = 2 | .remark = $remark' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit 40
            fi

            # --- 核心修改：根据 SOC_TYPE 下载对应的 Ops 包 ---
            case "$SOC_TYPE" in
                "A2") RUN_PKG_NAME="910b_ops.run" ;;
                "A3") RUN_PKG_NAME="A3_ops.run" ;;
                "A5") RUN_PKG_NAME="950_ops.run" ;;
            esac

            if ! wget -q -O "$RUN_DIR/$RUN_PKG_NAME" "$OPS_URL" --no-check-certificate --timeout=60 --tries=2; then
                jq --arg remark "41 download run pkg failed" \
                   '.status = 2 | .remark = $remark' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit 41
            fi

            chmod +x "$RUN_DIR"/*.run

            if [[ -f "$CANN_VERSION_INFO" ]]; then
                VERSION_VAL=$(grep "^Version=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
                TIMESTAMP_VAL=$(grep "^timestamp=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
                if [[ -n "$VERSION_VAL" && -n "$TIMESTAMP_VAL" ]]; then
                    mkdir -p "$(dirname "$CURRENT_VERSION_CACHE")" 2>/dev/null || true
                    cat > "$CURRENT_VERSION_CACHE" <<EOF
{
    "Version": "$VERSION_VAL",
    "timestamp": "$TIMESTAMP_VAL"
}
EOF
                fi
            fi

            ASCEND_ROOT="/home/developer/Ascend"
            ASCEND_BK="${ASCEND_ROOT}.bk"
            [[ -e "$ASCEND_BK" ]] && sudo rm -rf "$ASCEND_BK"
            if [[ -e "$ASCEND_ROOT" ]]; then
                if ! sudo mv "$ASCEND_ROOT" "$ASCEND_BK"; then
                    jq --arg remark "60 failed to backup Ascend directory" \
                       '.status = 2 | .remark = $remark' \
                       "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                    exit 60
                fi
            fi

            INSTALL_SUCCESS=true
            ERROR_CODE=0
            ERROR_MSG=""

            if ! (echo Y | "$RUN_DIR/toolkit.run" --install --force --install-path="$ASCEND_ROOT" >/dev/null 2>&1); then
                INSTALL_SUCCESS=false
                ERROR_CODE=50
                ERROR_MSG="50 install toolkit.run failed"
            # --- 核心修改：根据 SOC_TYPE 安装对应的 Ops 包 ---
            elif ! (echo Y | "$RUN_DIR/$RUN_PKG_NAME" --install --force --install-path="$ASCEND_ROOT" >/dev/null 2>&1); then
                INSTALL_SUCCESS=false
                ERROR_CODE=51
                ERROR_MSG="51 install ops.run failed"
            fi

            if [[ "$INSTALL_SUCCESS" == false ]]; then
                if [[ -e "$ASCEND_BK" ]]; then
                    sudo rm -rf "$ASCEND_ROOT"
                    if ! sudo mv "$ASCEND_BK" "$ASCEND_ROOT"; then
                        END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
                        cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "$START_TIME",
    "end_time": "$END_TIME",
    "process": "0%",
    "remark": "61 rollback failed after $ERROR_MSG"
}
EOF
                        exit 61
                    fi
                else
                    sudo rm -rf "$ASCEND_ROOT"
                fi

                END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
                jq --arg remark "$ERROR_MSG" \
                   --arg end_time "$END_TIME" \
                   '.status = 2 | .remark = $remark | .end_time = $end_time' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit "$ERROR_CODE"
            fi

            [[ -e "$ASCEND_BK" ]] && sudo rm -rf "$ASCEND_BK"
            END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
            jq --arg endtime "$END_TIME" \
               '.status = 1 | .end_time = $endtime | .process = "100%"' \
               "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
            rm -f "$CURRENT_VERSION_CACHE"
            source "$ASCEND_ROOT/cann/set_env.sh" 2>/dev/null || true
        ) &

        exit 0
        ;;
    "forceupgrade")
        force_cleanup_previous_upgrade >/dev/null 2>&1
        output_json '{"error_code":"0","error_msg":"Force upgrade task started after cleanup"}'

        (
            TEMP_DIR_ASYNC="/tmp/cann_upgrade_async_$$"
            mkdir -p "$TEMP_DIR_ASYNC" 2>/dev/null || true
            trap 'rm -rf "$TEMP_DIR_ASYNC"' EXIT

            # --- 同样：不用 sudo，先下载 ---
            if ! wget -q -O "$TEMP_DIR_ASYNC/versions.json" "$VERSIONS_URL" \
                    --no-check-certificate --timeout=20 --tries=2; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "30 download versions.json failed"
}
EOF
                exit 30
            fi

            exec >/dev/null 2>&1

            # --- 以下逻辑与 upgrade 完全一致 ---
            TARGET_ENTRY=$(jq -e --arg v "$TARGET_VERSION" '.versions[] | select(.version == $v)' "$TEMP_DIR_ASYNC/versions.json" 2>/dev/null) || true
            if [[ -z "$TARGET_ENTRY" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "20 target version not found"
}
EOF
                exit 20
            fi

            TOOLKIT_URL=$(echo "$TARGET_ENTRY" | jq -r ".download_url.${TARGET_ARCH}[\"Ascend-cann-toolkit_url\"]")
            # --- 核心修改：根据 SOC_TYPE 选择 Ops URL ---
            OPS_URL=$(echo "$TARGET_ENTRY" | jq -r ".download_url.${TARGET_ARCH}[\"$OPS_FIELD_NAME\"]")

            if [[ -z "$TOOLKIT_URL" || "$TOOLKIT_URL" == "null" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "22 toolkit URL missing for $TARGET_ARCH"
}
EOF
                exit 22
            fi
            if [[ -z "$OPS_URL" || "$OPS_URL" == "null" ]]; then
                cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "",
    "end_time": "",
    "process": "0%",
    "remark": "23 ops URL missing for $TARGET_ARCH"
}
EOF
                exit 23
            fi

            START_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
            cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 0,
    "target_version": "$TARGET_VERSION",
    "start_time": "$START_TIME",
    "end_time": "",
    "process": "0%",
    "remark": ""
}
EOF

            RUN_DIR="$TEMP_DIR_ASYNC/runs"
            mkdir -p "$RUN_DIR" 2>/dev/null || true
            if ! wget -q -O "$RUN_DIR/toolkit.run" "$TOOLKIT_URL" --no-check-certificate --timeout=60 --tries=2; then
                jq --arg remark "40 download run pkg failed" \
                   '.status = 2 | .remark = $remark' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit 40
            fi

            # --- 核心修改：根据 SOC_TYPE 下载对应的 Ops 包 ---
            case "$SOC_TYPE" in
                "A2") RUN_PKG_NAME="910b_ops.run" ;;
                "A3") RUN_PKG_NAME="A3_ops.run" ;;
                "A5") RUN_PKG_NAME="950_ops.run" ;;
            esac

            if ! wget -q -O "$RUN_DIR/$RUN_PKG_NAME" "$OPS_URL" --no-check-certificate --timeout=60 --tries=2; then
                jq --arg remark "41 download run pkg failed" \
                   '.status = 2 | .remark = $remark' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit 41
            fi

            chmod +x "$RUN_DIR"/*.run

            if [[ -f "$CANN_VERSION_INFO" ]]; then
                VERSION_VAL=$(grep "^Version=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
                TIMESTAMP_VAL=$(grep "^timestamp=" "$CANN_VERSION_INFO" | cut -d'=' -f2 | tr -d ' "')
                if [[ -n "$VERSION_VAL" && -n "$TIMESTAMP_VAL" ]]; then
                    mkdir -p "$(dirname "$CURRENT_VERSION_CACHE")" 2>/dev/null || true
                    cat > "$CURRENT_VERSION_CACHE" <<EOF
{
    "Version": "$VERSION_VAL",
    "timestamp": "$TIMESTAMP_VAL"
}
EOF
                fi
            fi

            ASCEND_ROOT="/home/developer/Ascend"
            ASCEND_BK="${ASCEND_ROOT}.bk"
            [[ -e "$ASCEND_BK" ]] && sudo rm -rf "$ASCEND_BK"
            if [[ -e "$ASCEND_ROOT" ]]; then
                if ! sudo mv "$ASCEND_ROOT" "$ASCEND_BK"; then
                    jq --arg remark "60 failed to backup Ascend directory" \
                       '.status = 2 | .remark = $remark' \
                       "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                    exit 60
                fi
            fi

            INSTALL_SUCCESS=true
            ERROR_CODE=0
            ERROR_MSG=""

            if ! (echo Y | "$RUN_DIR/toolkit.run" --install --force --install-path="$ASCEND_ROOT" >/dev/null 2>&1); then
                INSTALL_SUCCESS=false
                ERROR_CODE=50
                ERROR_MSG="50 install toolkit.run failed"
            # --- 核心修改：根据 SOC_TYPE 安装对应的 Ops 包 ---
            elif ! (echo Y | "$RUN_DIR/$RUN_PKG_NAME" --install --force --install-path="$ASCEND_ROOT" >/dev/null 2>&1); then
                INSTALL_SUCCESS=false
                ERROR_CODE=51
                ERROR_MSG="51 install ops.run failed"
            fi

            if [[ "$INSTALL_SUCCESS" == false ]]; then
                if [[ -e "$ASCEND_BK" ]]; then
                    sudo rm -rf "$ASCEND_ROOT"
                    if ! sudo mv "$ASCEND_BK" "$ASCEND_ROOT"; then
                        END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
                        cat > "$CANN_TASK_FILE" <<EOF
{
    "status": 2,
    "target_version": "$TARGET_VERSION",
    "start_time": "$START_TIME",
    "end_time": "$END_TIME",
    "process": "0%",
    "remark": "61 rollback failed after $ERROR_MSG"
}
EOF
                        exit 61
                    fi
                else
                    sudo rm -rf "$ASCEND_ROOT"
                fi

                END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
                jq --arg remark "$ERROR_MSG" \
                   --arg end_time "$END_TIME" \
                   '.status = 2 | .remark = $remark | .end_time = $end_time' \
                   "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
                exit "$ERROR_CODE"
            fi

            [[ -e "$ASCEND_BK" ]] && sudo rm -rf "$ASCEND_BK"
            END_TIME=$(date -u +"%Y%m%dT%H:%M:%SZ")
            jq --arg endtime "$END_TIME" \
               '.status = 1 | .end_time = $endtime | .process = "100%"' \
               "$CANN_TASK_FILE" > "${CANN_TASK_FILE}.tmp" && mv "${CANN_TASK_FILE}.tmp" "$CANN_TASK_FILE"
            rm -f "$CURRENT_VERSION_CACHE"
            source "$ASCEND_ROOT/cann/set_env.sh" 2>/dev/null || true
        ) &

        exit 0
        ;;
    *)
        fail 3 "Invalid action"
        ;;
esac

exit 0