#!/bin/bash
# 采集 NPU 状态并输出为 Influx Line Protocol

COMMON_OUTPUT=$(/usr/local/sbin/npu-smi info -t common -i 0 2>/dev/null)
MEM_OUTPUT=$(/usr/local/sbin/npu-smi info -t memory -i 0 2>/dev/null)

TEMP=$(echo "$COMMON_OUTPUT" | grep "Temperature(C)" | awk '{print $NF}')
UTIL=$(echo "$COMMON_OUTPUT" | grep "Aicore Usage Rate" | awk '{print $NF}')
MEM_TOTAL_MB=$(echo "$MEM_OUTPUT" | grep "Capacity(MB)" | awk '{print $NF}')
MEM_RATE=$(echo "$COMMON_OUTPUT" | grep "Memory Usage Rate" | awk '{print $NF}')

# 计算实际使用量 (使用 awk 代替 bc)
if [ -n "$MEM_TOTAL_MB" ] && [ -n "$MEM_RATE" ]; then
    MEM_USED_MB=$(awk "BEGIN {printf \"%d\", $MEM_TOTAL_MB * $MEM_RATE / 100}")
else
    MEM_TOTAL_MB=0
    MEM_USED_MB=0
    UTIL=0
    TEMP=0
fi

# 默认值保护
TEMP=${TEMP:-0}
UTIL=${UTIL:-0}

echo "npu_status,device=0 temp=${TEMP},utilization=${UTIL},memory_used=${MEM_USED_MB:-0},memory_total=${MEM_TOTAL_MB:-0}"