#!/usr/bin/env bash
# nvidia-oc [index] [busid]


TIMEOUT=15

[[ "$1" != "nocolor" ]] && source colors

if [[ "$1" == -h || "$1" == --help ]]; then
	echo "Usage: nvidia-info [--force] [index] [busid:00] ..."
	exit 0
fi


date


if [[ ! -f $GPU_DETECT_JSON ]]; then
	gpu_detect_json=`gpu-detect listjson`
else
	gpu_detect_json=$(< $GPU_DETECT_JSON)
fi


idx=0
while IFS=";" read busid brand name mem vbios fan_cnt; do
	BUSID[$idx]="$busid"
	BRAND[$idx]="$brand"
	NAME[$idx]="$name"
	RAM[$idx]="$mem"
	VBIOS[$idx]="$vbios"
	((idx++))
done < <( echo "$gpu_detect_json" | jq -r -c '.[] | select(.brand == "nvidia" or .vendor == "NVIDIA") | (.busid+";"+.brand+";"+.name+";"+.mem+";"+.vbios+";"+.fan_cnt)' 2>/dev/null )


if [[ ${#BUSID[@]} -eq 0 ]]; then
	echo "No NVIDIA cards detected"
	exit 0
fi

# get nvflash support
readarray -t NVFLASH < <( nvidia-flash --busid )


function nvidiainfo() {
	local nvdata=
	declare -A GPUINFO=()

	if [[ ! -f /run/hive/NV_OFF || " $@ " =~ " --force " ]]; then
		readarray -t nvdata < <(timeout -s9 $TIMEOUT nvtool -q --csv --busid --device --name --vbios --temp --fans --fancount --power \
									--curpl --defpl --minmaxpl --clocks --usage --pcie --meminfo --memtype --curgen --curwidth --sn --uuid \
									--state --throttle --memtemp --coreoffset --memoffset --clocksinfo --replay --maxtemp --statuscode --status)
	else
		echo "${YELLOW}NVIDIA driver is disabled. (use ${BYELLOW}--force${YELLOW} to override)${NOCOLOR}"
	fi
	if [[ ${#nvdata[@]} -eq 0 ]]; then
		echo -e "${RED}Driver is not loaded. Only basic info is provided${NOCOLOR}"
	else
		for gpuinfo in "${nvdata[@]}"; do
			# 00000000:0A:00.0;
			[[ ! "$gpuinfo" =~ ([0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.0) ]] && continue
			busid="${BASH_REMATCH[1],,}" # lower case
			GPUINFO["$busid"]="$gpuinfo"
		done
	fi

	for (( i=0; i < ${#BUSID[@]}; i++ )); do
		busid="${BUSID[$i]}"
		name="${NAME[$i]}"

		IFS=";" read bus_id index gpu_name bios temp rpm fan target fancnt pwr plcur pldef plmin plmax gclk mclk gpuutil memutil pcie \
			memtotal memused memfree memtype linkg linkw sn uuid pstate throttle mtemp coreofs memofs corelock memlock replay maxtemp status \
			msg etc < <( echo "${GPUINFO["$busid"]}" )

		# filter output by index or busid
		[[ ! -z "${@//--force}" && ! " $@ " =~ ( ${index:--} ) && ! " $@ " =~  ( ${busid/.0} ) ]] && continue

		echo ""

		if [[ "${BRAND[$i]}" != "nvidia" ]]; then
			echo "${YELLOW}===${NOCOLOR} GPU ${CYAN}-${NOCOLOR}, ${PURPLE}${busid} ${GREEN}$name${NOCOLOR} - ${BYELLOW}${VBIOS[$i]} ${YELLOW}===${NOCOLOR}"
			continue
		fi

		if [[ $status -ne 0 ]]; then
			echo "${YELLOW}===${NOCOLOR} GPU ${CYAN}-${NOCOLOR}, ${PURPLE}${busid} ${GREEN}${name}$NOCOLOR - ${RED}$msg ($status) ${YELLOW}===${NOCOLOR}"
			continue
		fi

		local faninfo=
		if [[ $fancnt -gt 0 ]]; then
			readarray -t fan_arr < <( echo "$fan" | tr "," "\n" )
			readarray -t rpm_arr < <( echo "$rpm" | tr "," "\n" )
			readarray -t target_arr < <( echo "$target" | tr "," "\n" )

			faninfo=", Fan $WHITE${target_arr[0]:--} %$NOCOLOR"
			for (( f=0; f < ${#rpm_arr[@]}; f++ )); do
				faninfo+=", RPM $WHITE${rpm_arr[f]}$NOCOLOR"
				[[ "${fan_arr[f]}" != "${target_arr[f]}" ]] && faninfo+=" (${fan_arr[f]} %)"
			done
		fi

		[[ ! -z "$pstate" ]] && pstate="P${pstate}"
		[[ ! -z "$linkw" ]] && linkw="${linkw}x"
		[[ "$throttle" =~ SW|HW ]] && throttle="${BRED}$throttle"
		[[ $mtemp -eq 0 ]] && mtemp= || mtemp=", Mem temp $WHITE${mtemp:--} °C$NOCOLOR"
		[[ $coreofs -eq 0 ]] && coreofs= || coreofs=" (${WHITE}`((coreofs > 0)) && echo "+"`$coreofs${NOCOLOR})"
		[[ $memofs -eq 0 ]] && memofs= || memofs=" (${WHITE}`((memofs > 0)) && echo "+"`$memofs${NOCOLOR})"
		[[ $corelock -eq 0 ]] && corelock= || corelock=" [${WHITE}$corelock${NOCOLOR}]"
		[[ $memlock -eq 0 ]] && memlock= || memlock=" [${WHITE}$memlock${NOCOLOR}]"
		[[ $replay -eq 0 ]] && replay= || replay=", Errors: ${WHITE}`((replay > 10)) && echo "${BRED}"`$replay${NOCOLOR}"
		[[ $maxtemp -eq 0 ]] && maxtemp= || maxtemp=" (max $WHITE${maxtemp:--} °C$NOCOLOR)"

		echo "${YELLOW}===${NOCOLOR} GPU ${CYAN}$index${NOCOLOR}, ${PURPLE}${busid} ${GREEN}${name} ${WHITE}$memtotal MB ${YELLOW}===${NOCOLOR}"
		echo "  Bios $WHITE${bios:--}$NOCOLOR, PCIE Link:  Gen $WHITE${linkg:--}$NOCOLOR, Width $WHITE${linkw:--}$NOCOLOR$replay"
		[[ ! -z "$serial" ]] && echo -n "  S/N ${serial:--}, UUID ${uuid:--}" || echo -n "  UUID ${uuid:--}"
		[[ " ${NVFLASH[@]} " =~ " ${busid/.0} " ]] && echo -n ", ROM flash: ${WHITE}Supported${NOCOLOR}" || echo -n ", ROM flash: ${WHITE}Not supported${NOCOLOR}"
		echo
		echo "  Power Limit:  Min $WHITE${plmin:--} W$NOCOLOR, Default $WHITE${pldef:--} W$NOCOLOR, Max $WHITE${plmax:--} W$NOCOLOR, Current $WHITE${plcur:--} W$NOCOLOR"
		echo "  Frequency:  Core $WHITE${gclk:--} MHz$NOCOLOR${corelock}${coreofs}, MEM $WHITE${mclk:--} MHz$NOCOLOR${memlock}${memofs}"
		echo "  Memory:  Total $WHITE${memtotal:--} MB$NOCOLOR, Used $WHITE${memused:--} MB$NOCOLOR, Free $WHITE${memfree:--} MB$NOCOLOR, $WHITE${memtype:--}$NOCOLOR"
		echo "  Utilization:  GPU $WHITE${gpuutil:--} %$NOCOLOR, MEM $WHITE${memutil:--} %$NOCOLOR, PCIE $WHITE${pcie:--} %$NOCOLOR, Throttle: $WHITE$throttle$NOCOLOR"
		echo "  PSTATE ${WHITE}${pstate:--}$NOCOLOR, PWR $WHITE${pwr:--} W$NOCOLOR, Temp $WHITE${temp:--} °C$NOCOLOR${maxtemp}${mtemp}${faninfo}"

	done
}


nvidiainfo $@
