#!/bin/bash

# ================= 配置区域 =================
OUTPUT_MODE="JSON_ONLY" # BOTH, TEXT_ONLY, JSON_ONLY
# ===========================================

# ================= 全局存储 =================
declare -a IPMI_ORDER=()
declare -a DMESG_ORDER=()
declare -a ERR_GPU_HIGH_TEMP=()
declare -a ERR_GPU_LOW_POWER=()
declare -a ERR_VBIOS=()
declare -a ERR_FAN=()
declare -a ERR_GPU_COUNT=()
declare -a ERR_LINK=()
declare -a ERR_BURN_FAIL=()
declare -a ERR_GPU_NVIDIA=()
declare -a ERR_MODEL=()

declare -A IPMI_COUNTS=()
declare -A DMESG_COUNTS=()
declare -A IPMI_DATA=()
declare -A DMESG_DATA=()
declare -A ERROR_TYPES=()

# ================= 基础信息初始化 =================
os_ver=$(. /etc/os-release && echo "$NAME")
SN=$(dmidecode -s system-serial-number 2>/dev/null | grep -v "^#" | head -1 | xargs)
MY_IP="${1:-$(hostname -I 2>/dev/null | awk '{print $1}')}"

[[  $os_ver =~ 'Ubuntu' ]] &&  cmdb_file='/mnt/tmp/os_init/cmdb.info' && eval `sed 's/ /\n/g' /proc/cmdline | grep gpu_task`  || cmdb_file='/tmp/os_init/cmdb.info'
[[  -f $cmdb_file ]] && IDC=`awk -F ':' '/IDC/{print $NF}' $cmdb_file`

if lspci -n 2>/dev/null | grep -q "10de:"; then
    HAS_GPU=true
    lspci -nn | grep -qw 10de:1eb8 && GPU_TYPE='T4'
    lspci -nn | grep -qw 10de:20f3 && GPU_TYPE='A800-SXM4-80GB'
    lspci -nn | grep -qw 10de:1e82 && GPU_TYPE='2080'
    lspci -nn | grep -qwE "10de:1e07|10de:1e04" && GPU_TYPE='2080Ti'
    lspci -nn | grep -qw 10de:2b85 && GPU_TYPE='5090'
    lspci -nn | grep -qw 10de:2208 && GPU_TYPE='3080'
    lspci -nn | grep -qw 10de:2329 && GPU_TYPE='H20'
    lspci -nn | grep -qw '10de:2684' && GPU_TYPE='4090'
    lspci -nn | grep -qw 10de:1db4 && GPU_TYPE='V100'
    lspci -nn | grep -qw '10de:1b38' && GPU_TYPE='P40'
else
    GPU_TYPE='CPU'
fi

stress_status=''
MEM_STRE_TIME=""
MODEL_RUN_TIME=""
uptime_min=0
check_time_iso=""
status_msg=""

# ================= 辅助函数 =================
json_escape_str() {
    local s="$1"
    if [[ -z "$s" ]]; then return; fi
    s="${s//\\/\\\\}"
    s="${s//\"/\\\"}"
    s="${s//$'\n'/\\n}"
    s="${s//$'\r'/\\r}"
    s="${s//$'\t'/\\t}"
    printf '%s' "$s"
}

get_ipmi_type() {
    local full_text="$*"
    if [[ "$full_text" == *[Tt]hermal* ]] || [[ "$full_text" == *[Tt]emp* ]] || [[ "$full_text" == *[Hh]ot* ]]; then echo "Thermal"
    elif [[ "$full_text" == *[Vv]oltage* ]] || [[ "$full_text" == *[Pp]ower* ]]; then echo "Voltage"
    elif [[ "$full_text" == *[Mm]emory* ]] || [[ "$full_text" == *ECC* ]]; then echo "Memory"
    elif [[ "$full_text" == *[Bb]us* ]] || [[ "$full_text" == *[Pp]CI* ]] || [[ "$full_text" == *[Ii]ERR* ]]; then echo "Bus Error"
    elif [[ "$full_text" == *[Ff]an* ]] || [[ "$full_text" == *[Ss]peed* ]]; then echo "Fan"
    elif [[ "$full_text" == *[Cc]ritical* ]] || [[ "$full_text" == *[Ff]ailure* ]] || [[ "$full_text" == *[Ff]ault* ]]; then echo "Critical"
    else echo "Fault"
    fi
}

add_ipmi_error() {
    local raw_line="$1"
    local parsed
    parsed=$(echo "$raw_line" | awk -F '|' '{
        n = NF
        for(i=1; i<=n; i++) { gsub(/^[ \t]+|[ \t]+$/, "", $i) }
        id = (n>=1) ? $1 : ""
        date_part = (n>=2) ? $2 : ""
        time_part = (n>=3) ? $3 : ""
        sensor = (n>=4) ? $4 : ""
        rest = ""
        for(i=5; i<=n; i++) { if($i != "") rest = (rest == "") ? $i : rest " " $i }
        printf "%s@%s@%s@%s@%s", id, date_part, time_part, sensor, rest
    }')

    local id date_part time_part sensor rest_content
    IFS='@' read -r id date_part time_part sensor rest_content <<< "$parsed"

    local event assertion
    if [[ -n "$rest_content" ]]; then
        assertion=$(echo "$rest_content" | awk '{print $NF}')
        event=$(echo "$rest_content" | sed "s/ ${assertion}$//")
        if [[ "$event" == "$assertion" ]]; then event=""; fi
    fi

    local err_type
    err_type=$(get_ipmi_type "$event" "$assertion")
    local fp="${sensor}|${event}"

    if [[ -z "${IPMI_COUNTS[$fp]}" ]]; then
        IPMI_COUNTS["$fp"]=1
        IPMI_DATA["$fp"]="$raw_line"
        IPMI_DATA["${fp}_ID"]="$id"
        IPMI_DATA["${fp}_DATE"]="$date_part"
        IPMI_DATA["${fp}_TIME"]="$time_part"
        IPMI_DATA["${fp}_SENSOR"]="$sensor"
        IPMI_DATA["${fp}_EVENT"]="$event"
        IPMI_DATA["${fp}_ASSERTION"]="$assertion"
        IPMI_DATA["${fp}_TYPE"]="$err_type"
        IPMI_ORDER+=("$fp")
    else
        ((IPMI_COUNTS["$fp"]++))
    fi
    ERROR_TYPES["IPMI_ERR"]=1
}

add_dmesg_error() {
    local raw_line="$1"
    local time_part="Unknown"
    local content_part=""
    local dtype="General Error"
    local event_desc=""

    if [[ "$raw_line" =~ ^\[([^\]]+)\][[:space:]]*(.*) ]]; then
        local raw_date="${BASH_REMATCH[1]}"
        content_part="${BASH_REMATCH[2]}"
        if command -v date &>/dev/null; then
            time_part=$(date -d "$raw_date" +"%Y-%m-%dT%H:%M:%S" 2>/dev/null || echo "$raw_date")
        else
            time_part="$raw_date"
        fi
    else
        content_part="$raw_line"
    fi

    local keyword=""
    local match_found=true

    if [[ "$content_part" == *"Xid"* ]]; then
        dtype="GPU Xid"; keyword="Xid"
    elif [[ "$content_part" == *"Hardware Error"* ]]; then
        dtype="Hardware Error"; keyword="Hardware Error"
    elif [[ "$content_part" == *"blocked for more than"* ]]; then
        dtype="Blocked"; keyword="blocked for more than"
    elif [[ "$content_part" == *"out of memory"* ]]; then
        dtype="OOM"; keyword="out of memory"
    elif [[ "$content_part" == *"Call Trace"* ]]; then
        dtype="Call Trace"; event_desc="Kernel call trace detected"
    elif [[ "$content_part" == *"Temperature above threshold"* ]]; then
        dtype="Thermal"; keyword="Temperature above threshold"
    else
        match_found=false
    fi

    if [[ "$match_found" == true && -n "$keyword" && "$dtype" != "Call Trace" ]]; then
        local suffix="${content_part#*$keyword}"
        event_desc=$(echo "$suffix" | sed 's/^[^a-zA-Z0-9]*//')
        if [[ -z "$event_desc" ]]; then
            event_desc="${content_part:0:60}"
        fi
    elif [[ "$match_found" == false ]]; then
        event_desc="${content_part:0:60}"
    fi
    
    event_desc=$(echo "$event_desc" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
    if [[ -z "$event_desc" ]]; then event_desc="Unknown Event"; fi

    local content_key="${dtype}:${event_desc:0:40}"

    if [[ -z "${DMESG_COUNTS[$content_key]}" ]]; then
        DMESG_COUNTS["$content_key"]=1
        DMESG_DATA["$content_key"]="${time_part}|${dtype}|${event_desc}|${raw_line}"
        DMESG_ORDER+=("$content_key")
    else
        ((DMESG_COUNTS["$content_key"]++))
    fi
    ERROR_TYPES["DMESG_ERR"]=1
}

get_error_severity() {
    local code="$1"
    local subtype="$2"

    case "$code" in
        BURN_FAIL_ERR|LINK_ERR|GPU_NVIDIA_ERR|GPU_COUNT_ERR)
            echo "critical"
            ;;
        GPU_HIGH_TEMP_ERR|GPU_LOW_POWER_ERR|FAN_ERR|VBIOS_ERR|MODEL_ERR)
            echo "warning"
            ;;
        IPMI_ERR)
            case "$subtype" in
                Critical|Bus\ Error|Memory)
                    echo "critical"
                    ;;
                *)
                    echo "warning"
                    ;;
            esac
            ;;
        DMESG_ERR)
            case "$subtype" in
                GPU\ Xid|Hardware\ Error|Call\ Trace)
                    echo "critical"
                    ;;
                *)
                    echo "warning"
                    ;;
            esac
            ;;
        *)
            echo "warning"
            ;;
    esac
}

build_error_summary() {
    local target="$1"
    local condition="$2"
    local message="$3"
    local raw="$4"
    local summary=""

    if [[ -n "$condition" && -n "$target" ]]; then
        summary="$target $condition"
    elif [[ -n "$condition" ]]; then
        summary="$condition"
    elif [[ -n "$message" ]]; then
        summary="$message"
    elif [[ -n "$raw" ]]; then
        summary="$raw"
    fi

    if [[ -n "$message" && "$message" != "$summary" ]]; then
        if [[ -n "$summary" ]]; then
            summary="$summary - $message"
        else
            summary="$message"
        fi
    fi

    printf '%s' "$summary"
}

add_struct_err() {
    local arr_name="$1"
    shift
    local json_obj="{"
    local first=true
    local args=("$@")
    local total_args=${#args[@]}
    local i=0
    local target=""
    local raw_val=""
    local message_val=""
    local condition_val=""

    while [[ $i -lt $total_args ]]; do
        local k="${args[$i]}"
        local v="${args[$((i+1))]}"
        if [[ -n "$v" && "$v" != "N/A" && "$v" != "null" ]]; then
            case "$k" in
                bus_id|gpu_bus_id|slot|sensor|task|fault_bus_ids)
                    [[ -z "$target" ]] && target="$v"
                    ;;
            esac
            [[ "$k" == "raw" ]] && raw_val="$v"
            [[ "$k" == "message" ]] && message_val="$v"
            [[ "$k" == "condition" ]] && condition_val="$v"
            if [[ "$first" == true ]]; then first=false; else json_obj+=","; fi
            if [[ "$v" =~ ^-?[0-9]+\.?[0-9]*$ ]]; then
                json_obj+="\"$k\":$v"
            else
                json_obj+="\"$k\":\"$(json_escape_str "$v")\""
            fi
        fi
        ((i+=2))
    done

    local code="${arr_name#ERR_}_ERR"
    local severity
    severity=$(get_error_severity "$code" "")
    local summary
    summary=$(build_error_summary "$target" "$condition_val" "$message_val" "$raw_val")

    if [[ -n "$code" ]]; then
        if [[ "$first" == true ]]; then first=false; else json_obj+=","; fi
        json_obj+="\"code\":\"$(json_escape_str "$code")\""
    fi
    if [[ -n "$severity" ]]; then
        if [[ "$first" == true ]]; then first=false; else json_obj+=","; fi
        json_obj+="\"severity\":\"$(json_escape_str "$severity")\""
    fi
    if [[ -n "$target" ]]; then
        if [[ "$first" == true ]]; then first=false; else json_obj+=","; fi
        json_obj+="\"target\":\"$(json_escape_str "$target")\""
    fi
    if [[ -n "$summary" ]]; then
        if [[ "$first" == true ]]; then first=false; else json_obj+=","; fi
        json_obj+="\"summary\":\"$(json_escape_str "$summary")\""
    fi

    json_obj+="}"
    eval "$arr_name+=('$json_obj')"
    local type_flag="${arr_name#ERR_}"
    ERROR_TYPES["$type_flag"]=1
}

count_error_sections() {
    local total=0

    [[ ${#IPMI_ORDER[@]} -gt 0 ]] && ((total++))
    [[ ${#DMESG_ORDER[@]} -gt 0 ]] && ((total++))

    for type_flag in "${!ERROR_TYPES[@]}"; do
        if [[ "$type_flag" == "IPMI_ERR" || "$type_flag" == "DMESG_ERR" ]]; then continue; fi
        local arr_name="ERR_$type_flag"
        local count=0
        eval "count=\${#${arr_name}[@]}"
        [[ $count -gt 0 ]] && ((total++))
    done

    printf '%s' "$total"
}

count_error_items() {
    local total=0

    ((total += ${#IPMI_ORDER[@]}))
    ((total += ${#DMESG_ORDER[@]}))

    for type_flag in "${!ERROR_TYPES[@]}"; do
        if [[ "$type_flag" == "IPMI_ERR" || "$type_flag" == "DMESG_ERR" ]]; then continue; fi
        local arr_name="ERR_$type_flag"
        local count=0
        eval "count=\${#${arr_name}[@]}"
        ((total += count))
    done

    printf '%s' "$total"
}

status_has_issue() {
    local is_bad=false

    if [[ -n "$stress_status" && "$stress_status" != "mem_stress_end" && "$stress_status" != "model_stress_end" ]]; then
        is_bad=true
    fi
    if [[ ! "$os_ver" =~ 'Ubuntu' ]]; then
        is_bad=true
    fi

    printf '%s' "$is_bad"
}

# ================= 检查执行函数 =================

run_common_checks() {
    local ipmi_filter="Critical|IERR|Thermal Trip|Non-recoverable|Uncorrectable|Correctable|Fault|Failure|Device Disabled|ECC|error|input lost|Degraded|bus error|Bus Fatal Error|AC lost|over current|over voltage|memory ECC"
    local ipmi_exclude="OS Critical Stop OS Status"

    while IFS= read -r line; do
        [[ -n "$line" ]] && add_ipmi_error "$line"
    done < <(timeout 60 ipmitool sel elist 2>/dev/null | grep -iE "$ipmi_filter" | grep -v "$ipmi_exclude" | tac)

    if lspci -n |grep -q "10de:";then
        local dmesg_filter="call trace|Hardware Error|Xid|out of memory|soft lockup|over current|scsi hang|Blocked for more than|machine check events|resetting link|Media error|I/O error|task abort|correctable error|Ata error|Unrecovered read error|Temperature above threshold|transmit timed out|RAS.ERROR|ATU.ERROR|MON.ERROR|RmInitAdapter failed"
    else
        local dmesg_filter="Hardware Error|Xid|out of memory|soft lockup|over current|scsi hang|Blocked for more than|machine check events|resetting link|Media error|I/O error|task abort|correctable error|Ata error|Unrecovered read error|Temperature above threshold|transmit timed out|RAS.ERROR|ATU.ERROR|MON.ERROR|RmInitAdapter failed"
    fi

    while IFS= read -r line; do
        [[ -n "$line" ]] && add_dmesg_error "$line"
    done < <(dmesg -T 2>/dev/null | sed -n '/init as init process/,$p' | grep -iE "$dmesg_filter" | tail -n 10 | tac)
    # done 已确认存在
}

run_gpu_checks() {

    local smi_raw_output
    smi_raw_output=$(timeout 10 nvidia-smi --query-gpu=index,pci.bus_id,fan.speed --format=csv,noheader 2>&1)
    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        add_struct_err "ERR_GPU_NVIDIA" "condition" "nvidia-smi execution failed" "message" "$smi_raw_output"
        return
    fi

    if [[ "$GPU_TYPE" == *"A800"* ]]; then
        local v_bios
        v_bios=$(nvidia-smi -q 2>/dev/null | awk '/VBIOS/{print $NF}' | sort -u)
        if [[ -n "$v_bios" && "$v_bios" != "92.00.AC.00.03" ]]; then
            add_struct_err "ERR_VBIOS" "condition" "v_bios mismatch" "current_version" "$v_bios" "expected_version" "92.00.AC.00.03" "message" "v_bios is $v_bios, not 92.00.AC.00.03"
        fi
    fi

    mapfile -t power <<< "$smi_raw_output"
    if [[ ${#power[@]} -eq 1 && -z "${power[0]}" ]]; then power=(); fi

    if [[ "$GPU_TYPE" != "T4" && "$GPU_TYPE" != "A800-SXM4-80GB" && "$GPU_TYPE" != "H20" && "$GPU_TYPE" != "V100" && "$GPU_TYPE" != "P40" ]]; then
        local fan_speed
       # fan_speed=$(printf "%s\n" "${power[@]}" | awk -F ',' '{ if(substr($NF, length($NF)) != "%") print $(NF-1)}'  | sort -u |sed 's/00000000://g')
        fan_speed=$(printf "%s\n" "${power[@]}" | awk -F ',' '{ 
            val=$NF; gsub(/^[ \t]+|[ \t]+$/, "", val); 
            if(val !~ /%$/) { v=$(NF-1); gsub(/^[ \t]+|[ \t]+$/, "", v); print v } 
        }' | sort -u | sed 's/00000000://g')
        if [[ -n "$fan_speed" ]]; then
            add_struct_err "ERR_FAN" "condition" "fan_speed format error" "gpu_bus_id" "$fan_speed" "message" "fan_speed value missing '%' suffix"
        fi
    fi

    local num_elements="${#power[@]}"
    if [[ "$num_elements" -ne 8 && "$num_elements" -ne 4 && "$num_elements" -ne 10 ]]; then
        add_struct_err "ERR_GPU_COUNT" "condition" "count mismatch" "count" "$num_elements" "expected" "4 or 8" "message" "NVIDIA GPU count is $num_elements"
    fi

    if [[ -f /tmp/gpu-temp-pwr.csv ]]; then
        local b_temp=85
        local temp_details
        temp_details=$(awk -F', *' -v b_t="$b_temp" '
        /gpu_burn$/{
            bus=$(NF-2); gsub(/00000000:/,"",bus)
            total[bus]++; tmp=$6+0
            if(tmp>b_t){ over[bus]++; if(tmp>maxt[bus]) maxt[bus]=tmp; last_over[bus]=total[bus] }
        }
        END{ for(b in over) printf "%s|%d|%d|%d|%.1f|%d\n", b, maxt[b], total[b], over[b], (over[b]*100/total[b]), last_over[b] }
        ' /tmp/gpu-temp-pwr.csv)

        if [[ -n "$temp_details" ]]; then
            while IFS='|' read -r bus max_t total_s over_c pct last_idx; do
                local condition="MAX_temp>${b_temp}"
                local raw_str="${condition}, gpu_bus_id: ${bus} max:${max_t} (num:${over_c}/total:${total_s}):${pct}% last_num:${last_idx}"
                add_struct_err "ERR_GPU_HIGH_TEMP" "condition" "$condition" "bus_id" "$bus" "max_temp" "$max_t" "total_samples" "$total_s" "over_count" "$over_c" "percentage" "$pct" "last_over_index" "$last_idx" "raw" "$raw_str"
            done <<< "$temp_details"
        fi
    fi

    if command -v nvidia-smi &> /dev/null; then
        mapfile -t arr < <( nvidia-smi --query-gpu=power.limit --format=csv,noheader,nounits 2>/dev/null | sort -u | awk '{print int($0)}')
        [[  ${#arr[@]} -gt 1 ]] && add_struct_err "ERR_GPU_power_limit_unique"  "messages" "GPU power.limit:${arr[*]} not unique, check max via --query-gpu=power.max_limit"
        b_power=$(awk "BEGIN{print int(${arr[0]} * 0.92)}")
        local calc_str="${arr[0]}*0.92"
            
        if [[ -f /tmp/gpu-temp-pwr.csv ]]; then
            local low_details
            low_details=$(awk -F', *' -v bp="$b_power"  '/gpu_burn$/{
            bus=$(NF-2); gsub(/00000000:/,"",bus); total[bus]++
            if(match($0, /([0-9.]+)[[:space:]]*W/, a)) w=a[1]+0; else next
            if(++seen[bus] <= 5) next
            if(w+0<bp){ low[bus]++; if(min_all[bus]==""||w<min_all[bus]) min_all[bus]=w; if(max_all[bus]==""||w>max_all[bus]) max_all[bus]=w; last_low[bus]=total[bus] }
            }
            END{ for(b in low) printf "%s|%d|%d|%d|%d|%.1f|%d\n", b, min_all[b], max_all[b], total[b], low[b], (low[b]*100/total[b]), last_low[b] }
            ' /tmp/gpu-temp-pwr.csv)

                if [[ -n "$low_details" ]]; then
                    while IFS='|' read -r bus min_w max_w total_s low_c pct last_idx; do
                        local condition="Power < ${b_power}W (${calc_str})"
                        local raw_str="(Power)<${b_power}w (${calc_str}), gpu_bus_id: ${bus} ${min_w}-${max_w}W (num:${low_c}/total:${total_s}):${pct}% last_num:${last_idx}"
                        add_struct_err "ERR_GPU_LOW_POWER" "condition" "$condition" "bus_id" "$bus" "min_watt" "$min_w" "max_watt" "$max_w" "total_samples" "$total_s" "low_count" "$low_c" "percentage" "$pct" "last_low_index" "$last_idx" "limit_watt" "$b_power" "raw" "$raw_str"
                    done <<< "$low_details"
                fi
            fi
    fi

}

run_link_check() {
    local slots
    slots=$(lspci 2>/dev/null | grep -iEw "vga|3D" | grep -i nvidia | awk '{print $1}')
    for slot in $slots; do
        local LnkSta last_word
        LnkSta=$(lspci -vvv -s "$slot" 2>/dev/null | grep 'LnkSta:')
        if [[ -n "$LnkSta" ]]; then
            last_word=$(echo "$LnkSta" | awk '{print $NF}')
            if [[ "$last_word" != *[Oo][Kk]* ]]; then
                add_struct_err "ERR_LINK" "condition" "Lnk_err" "slot" "$slot" "raw_data" "$LnkSta" "message" "Lnk_err"
            fi
        fi
    done
}

run_burn_check() {
    local gpu_num gpu_OK_num gpu_fault_id bus_ids
    gpu_num=$(nvidia-smi -L 2>/dev/null | wc -l)
    gpu_OK_num=0
    gpu_fault_id=""
    
    if [[ -f /tmp/gpu-burn1.log ]]; then
        local burn_content
        burn_content=$(tail -n 10 /tmp/gpu-burn1.log | grep -A 10 "GPUs:" | grep -E 'GPU [0-9]' || true)
        if [[ -n "$burn_content" ]]; then
            # grep -c prints 0 and exits 1 when there are no matches; do not use "|| echo 0"
            # or the capture becomes two lines and breaks numeric comparisons.
            gpu_OK_num=$(echo "$burn_content" | grep -c "OK" || true)
            gpu_fault_id=$(echo "$burn_content" | grep -v "OK" | awk '{print $2}' | sed 's/://g')
        fi
    fi

    if [[ $gpu_OK_num -lt 1 || -n "$gpu_fault_id" ]]; then
        bus_ids=""
        for index in $gpu_fault_id; do
            local bid
            bid=$(nvidia-smi --query-gpu=index,gpu_bus_id --format=csv,noheader 2>/dev/null | awk -F',' -v idx="$index" '$1==idx{print $2}' | sed 's/00000000://g')
            if [[ -n "$bid" ]]; then bus_ids="$bus_ids $bid"; fi
        done
        add_struct_err "ERR_BURN_FAIL" "condition" "Burn test failed" "total_gpus" "$gpu_num" "ok_gpus" "$gpu_OK_num" "fault_bus_ids" "$(echo $bus_ids | xargs)" "message" "G_num:$gpu_num G_OK_num:$gpu_OK_num"
    fi
}

run_model_check() {
    if [[ "$gpu_task" != "all" ]]; then return; fi
    if [[ "$stress_status" != "mem_stress_end" ]]; then return; fi

    local model_pro=""
    if [[ "$GPU_TYPE" == '5090' ]]; then 
        model_pro=$(pgrep VLLM | head -n 1)
    else
        model_pro=$(ps -ef | awk '/RayWor/ && !/awk/{print $3}' | sort -u | head -n 1)
    fi

    if [[ -z "$model_pro" ]]; then
        stress_status='GPU_MODEL_NOT_RUN'
        add_struct_err "ERR_MODEL" "condition" "Model process not found" "task" "$gpu_task" "message" "GPU Model (Ray/VLLM) not running"
    else
        local e_times
        e_times=$(ps -p "$model_pro" -o etimes= 2>/dev/null | tr -d ' ')
        if [[ -n "$e_times" && "$e_times" -gt 0 ]]; then
            MODEL_RUN_TIME="$((e_times/3600)):$((e_times%3600/60))"
            if [[ $e_times -lt 43200 ]]; then 
                stress_status='model_stressing'
            else 
                stress_status='model_stress_end'
            fi
        fi
    fi
}

# ================= 输出生成函数 =================

generate_text_output() {
    local base_core="$SN $MY_IP IDC:$IDC type:$GPU_TYPE"
    [[ -n "$gpu_task" ]] &&  base_core="$base_core task:$gpu_task"

    local status_suffix=""
    if [[ -n "$stress_status" ]]; then status_suffix="$stress_status"; fi
    if [[ -n "$status_msg" ]]; then status_suffix="$status_msg"; fi

    if [[ -n "$status_suffix" ]]; then
        echo "$base_core $status_suffix"
    else
        echo "$base_core"
    fi

    if [[ ${#IPMI_ORDER[@]} -gt 0 ]]; then
        echo "IPMI_ERR:"
        for fp in "${IPMI_ORDER[@]}"; do echo "   ${IPMI_DATA[$fp]}"; done
    fi
    
    if [[ ${#DMESG_ORDER[@]} -gt 0 ]]; then
        echo "DMESG_ERR:"
        for fp in "${DMESG_ORDER[@]}"; do 
            local data="${DMESG_DATA[$fp]}"
            local raw_line
            raw_line=$(echo "$data" | cut -d'|' -f4-)
            echo "   $raw_line"
        done
    fi
    
    for type_flag in "${!ERROR_TYPES[@]}"; do
        if [[ "$type_flag" == "IPMI_ERR" || "$type_flag" == "DMESG_ERR" ]]; then continue; fi
        local arr_name="ERR_$type_flag"
        local count=0
        eval "count=\${#${arr_name}[@]}"
        if [[ $count -gt 0 ]]; then
            echo "${type_flag}:"
            local i=0
            while [[ $i -lt $count ]]; do
                local item
                eval "item=\${${arr_name}[$i]}"
                local msg
                msg=$(echo "$item" | grep -o '"raw":"[^"]*"' | cut -d'"' -f4)
                if [[ -z "$msg" ]]; then msg=$(echo "$item" | grep -o '"message":"[^"]*"' | cut -d'"' -f4); fi
                msg=$(echo -e "$msg")
                echo "   $msg"
                ((i++))
            done
        fi
    done
}

generate_json_output() {
    local total_err total_items
    total_err=$(count_error_sections)
    total_items=$(count_error_items)
    echo "{"
    local type_val="CPU"
    if [[ "$HAS_GPU" == true ]]; then type_val="$GPU_TYPE"; fi
    
    local status_val="$stress_status"
    local b_json
    b_json="\"sn\":\"$(json_escape_str "$SN")\",\"ip\":\"$(json_escape_str "$MY_IP")\",\"idc\":\"$(json_escape_str "$IDC")\",\"type\":\"$(json_escape_str "$type_val")\",\"check_time\":\"$(json_escape_str "$check_time_iso")\",\"status\":\"$(json_escape_str "$status_val")\""
    
    if [[ -n "$gpu_task" ]]; then b_json+=",\"gpu_task\":\"$(json_escape_str "$gpu_task")\""; fi
    if [[ -n "$MEM_STRE_TIME" ]]; then b_json+=",\"mem_stre_time\":\"$(json_escape_str "$MEM_STRE_TIME")\""; fi
    if [[ -n "$MODEL_RUN_TIME" ]]; then b_json+=",\"model_run_time\":\"$(json_escape_str "$MODEL_RUN_TIME")\""; fi
    if [[ -n "$status_msg" ]]; then b_json+=",\"status_msg\":\"$(json_escape_str "$status_msg")\""; fi
    
    echo "  \"base_info\": {$b_json},"
    echo "  \"summary\": {\"error_count\":$total_err,\"error_type_count\":$total_err,\"error_item_count\":$total_items},"

    # 收集所有要输出的错误类型
    local -a output_sections=()
    local -a types_arr=()

    # --- IPMI_ERR ---
    if [[ ${#IPMI_ORDER[@]} -gt 0 ]]; then
        local -a ipmi_items=()
        for fp in "${IPMI_ORDER[@]}"; do
            local raw="${IPMI_DATA[$fp]}"
            local item
            item="\"$(json_escape_str "$raw")\""
            ipmi_items+=("$item")
        done
        
        local joined
        joined=$(IFS=','; echo "${ipmi_items[*]}")
        output_sections+=("    \"IPMI_ERR\": [$joined]")
        types_arr+=("\"IPMI_ERR\"")
    fi

    # --- DMESG_ERR ---
    if [[ ${#DMESG_ORDER[@]} -gt 0 ]]; then
        local -a dmesg_items=()
        for fp in "${DMESG_ORDER[@]}"; do
            local data="${DMESG_DATA[$fp]}"
            local d_raw
            d_raw=$(echo "$data" | cut -d'|' -f4-)
            
            local item
            item="\"$(json_escape_str "$d_raw")\""
            dmesg_items+=("$item")
        done
        
        local joined
        joined=$(IFS=','; echo "${dmesg_items[*]}")
        output_sections+=("    \"DMESG_ERR\": [$joined]")
        types_arr+=("\"DMESG_ERR\"")
    fi

    # --- Other Errors ---
    for type_flag in "${!ERROR_TYPES[@]}"; do
        if [[ "$type_flag" == "IPMI_ERR" || "$type_flag" == "DMESG_ERR" ]]; then continue; fi
        local arr_name="ERR_$type_flag"
        local count=0
        eval "count=\${#${arr_name}[@]}"
        
        if [[ $count -gt 0 ]]; then
            local -a struct_items=()
            for ((i=0; i<count; i++)); do
                local item; eval "item=\${${arr_name}[$i]}"
                local raw_val
                raw_val=$(echo "$item" | grep -o '"raw":"[^"]*"' | cut -d'"' -f4)
                if [[ -z "$raw_val" ]]; then raw_val=$(echo "$item" | grep -o '"raw_data":"[^"]*"' | cut -d'"' -f4); fi
                if [[ -z "$raw_val" ]]; then raw_val=$(echo "$item" | grep -o '"message":"[^"]*"' | cut -d'"' -f4); fi
                if [[ -z "$raw_val" ]]; then raw_val=$(echo "$item" | grep -o '"condition":"[^"]*"' | cut -d'"' -f4); fi
                struct_items+=("\"$(json_escape_str "$raw_val")\"")
            done
            
            local joined
            joined=$(IFS=','; echo "${struct_items[*]}")
            local json_key="${type_flag}_ERR"
            output_sections+=("    \"$json_key\": [$joined]")
            types_arr+=("\"$json_key\"")
        fi
    done

    local types_str=""
    if [[ ${#types_arr[@]} -gt 0 ]]; then
        types_str=$(IFS=,; echo "${types_arr[*]}")
    fi
    echo "  \"error_types\": [$types_str],"

    echo "  \"errors\": {"

    # 输出所有 section，正确处理逗号
    local total_sections=${#output_sections[@]}
    local idx=0
    for section in "${output_sections[@]}"; do
        echo -n "$section"
        if [[ $idx -lt $((total_sections - 1)) ]]; then
            echo ","
        else
            echo ""
        fi
        ((idx++))
    done
    echo "  }"

    echo "}"
}

# ================= 主流程 =================

if [[ "$os_ver" =~ 'Ubuntu' ]]; then
   [[ -z "$GPU_TYPE" ]]&& GPU_TYPE=$(nvidia-smi -L 2>/dev/null | awk -F':' '{print $2}' | awk '{print $(NF-1)}' | uniq | head -1)

    mem_pro=$(pidof stressapptest)
    uptime_min=$(awk '{print int($1/60)}' /proc/uptime)

    if [[ -n "$mem_pro" ]]; then
        stress_status='mem_stress_ing'
        sec=$(awk '{print int($1)}' /proc/uptime)
        MEM_STRE_TIME=$(printf "%02d:%02d:%02d" "$((sec/3600))" "$(((sec%3600)/60))" "$((sec%60))")
    elif [[ $uptime_min -lt 25 && -z "$mem_pro" ]]; then
        stress_status='STRESS_NOT_STARTED'
        status_msg="uptime:${uptime_min}min, maybe stress not start"
    else
        stress_status='mem_stress_end'
    fi

    run_common_checks

    if [[ "$HAS_GPU" == true ]]; then
        run_gpu_checks
        run_link_check
        if [[ -z "$(pidof stressapptest)" ]]; then 
            run_burn_check
            run_model_check
        fi
    fi

else
    stress_status="not_ubuntu_env"
    status_msg="not Ubuntu"
    run_common_checks
fi

check_time_iso=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

if [[ "$OUTPUT_MODE" == "BOTH" || "$OUTPUT_MODE" == "TEXT_ONLY" ]]; then 
    generate_text_output
fi

if [[ "$OUTPUT_MODE" == "BOTH" || "$OUTPUT_MODE" == "JSON_ONLY" ]]; then 
    generate_json_output
fi

exit 0
