#!/bin/bash

. colors

ERR_OTHER=10
ERR_NO_GPU=11
ERR_INVALID_ARG=12
ERR_FILE_NOT_FOUND=13

force=0

n=`gpu-detect NVIDIA`

if [[ $n == 0 ]]; then
	echo "No NVIDIA GPUs detected"
	exit $ERR_OTHER
fi


declare -A IDX=()
declare -a BUSID=()

get_busid() {
	#NVIDIA display adapters present in system:
	#<0> GeForce GTX 1080     (10DE,1B80,10DE,1B80) S:00,B:01,D:00,F:00
	readarray -t nvdata < <( nvflash_linux --list )
	for line in "${nvdata[@]}"; do
		if [[ "$line" =~ \<([0-9]+)\>.*B:([0-9a-fA-F]+) ]]; then
			local busid="${BASH_REMATCH[2],,}:00"
			local index="${BASH_REMATCH[1]}"
			IDX["$busid"]="$index"
			BUSID[$index]="$busid"
		fi
	done
}


print_busid() {
	local busid
	for busid in "${BUSID[@]}"; do
		echo "$busid"
	done
}


gpu_index=
get_index() {
	gpu_index=
	local gpu="$1"
	if [[ "${gpu/.0}" =~ ^([0-9a-fA-F]{2}+:00)$ ]]; then
		local busid="${BASH_REMATCH[1]}"
		if [[ -z "${IDX["$busid"]}" ]]; then
			echo "${RED}ERROR: GPU ${BRED}$gpu${RED} not found or not supported${NOCOLOR}"
			return $ERR_NO_GPU
		fi
		gpu="${IDX["$busid"]}"
	elif [[ ! "$gpu" =~ ^[0-9]+$ ]]; then
		echo "${RED}ERROR: Incorrect GPU specified - ${BRED}$gpu${NOCOLOR}"
		return $ERR_INVALID_ARG
	elif [[ -z "${BUSID[gpu]}" ]]; then
		echo "${RED}ERROR: GPU ${BRED}$gpu${RED} not found or not supported${NOCOLOR}"
		return $ERR_NO_GPU
	fi
	gpu_index=$((gpu + 0))
}


cleanup() {
	# cleanup nvflash output
	echo "${1//$'\r'$'\n'/$'\n'}" | grep -vP "\x0d" | cat -s
}


save_rom() {
	if [[ $# -lt 2 ]]; then
		echo "${RED}ERROR: GPU or file not specified${NOCOLOR}"
		return $ERR_INVALID_ARG
	fi

	local filename="${@:${#@}}" # last param

	if [[ -z "$filename" ]]; then
		echo "${RED}ERROR: File not specified${NOCOLOR}"
		return $ERR_INVALID_ARG
	fi

	get_index "$1" || return

	echo
	echo "${YELLOW}===${NOCOLOR} GPU ${CYAN}$gpu_index${NOCOLOR}, ${PURPLE}${BUSID[$gpu_index]} ${WHITE}DOWNLOADING VBIOS ${YELLOW}=== $(date +%T)${NOCOLOR}"
	local exitcode
	echo "nvflash_linux -i$gpu_index -b $filename"
	nvflash_linux -i$gpu_index -b "$filename"
	exitcode=$?
	[[ $exitcode -eq 0 && ! -f "$filename" ]] && exitcode=$ERR_FILE_NOT_FOUND
	return $exitcode
}


save_all() {
	local ret=0
	local index
	local exitcode

	local filename="${@:${#@}}" # last param

	if [[ -z "$filename" ]]; then
		echo "${RED}ERROR: File not specified${NOCOLOR}"
		return $ERR_INVALID_ARG
	fi

	# split filename
	full="${filename##*/}"
	path="${filename:0:${#filename} - ${#full}}"
	name="${full%%.*}"
	ext="${full:${#name} + 1}"

	for index in "${!BUSID[@]}"; do
		# /path/name[GPU].ext
		save_rom $index "${path}${name}${index}${ext:+.}${ext}"
		exitcode=$?
		[[ $exitcode -ne 0 ]] && ret=$exitcode
	done
	return $ret
}


forced_flash() {
	local gpu_index="$1"
	local rom="$2"
	expect -c '
exp_internal 0
set timeout 120
log_user 0
spawn nvflash_linux -s -A -i'$gpu_index' -6 "'$rom'" 2>&1
log_user 1
expect {
    "*confirm*" { send -- "y" }
    eof { catch wait result; exit [lindex $result 3] }
    }
expect eof
catch wait result
exit [lindex $result 3]
'
}


flash_rom() {
	if [[ $# -lt 2 ]]; then
		echo "${RED}ERROR: No GPU or file specified${NOCOLOR}"
		return $ERR_INVALID_ARG
	fi

	local filename="${@:${#@}}" # last param

	if [[ -z "$filename" || ! -f "$filename" ]]; then
		echo "${RED}ERROR: File not found - ${BRED}$filename${NOCOLOR}"
		return $ERR_FILE_NOT_FOUND
	fi

	get_index "$1" || return

	local extra_args=
	for (( i=2; i<$#; i++ )); do
		extra_args="$extra_args ${!i}"
	done
	[[ " $extra_args " =~ " --force " ]] && force=1

	echo
	echo "${YELLOW}===${NOCOLOR} GPU ${CYAN}$gpu_index${NOCOLOR}, ${PURPLE}${BUSID[$gpu_index]} ${WHITE}FLASHING VBIOS ${YELLOW}=== $(date +%T)${NOCOLOR}"
	local exitcode
	nvflash_linux -s -i$gpu_index -r 2>&1
	if [[ $force -eq 1 ]]; then
		forced_flash $gpu_index "$filename"
	else
		nvflash_linux -s -A -i$gpu_index "$filename" 2>&1
	fi
	exitcode=$?
	return $exitcode
}


flash_all() {
	local ret=0
	local index
	local exitcode

	local filename="${@:${#@}}" # last param

	if [[ -z "$filename" || ! -f "$filename" ]]; then
		echo "${RED}ERROR: File not found - ${BRED}$filename${NOCOLOR}"
		return $ERR_FILE_NOT_FOUND
	fi

	for index in "${!BUSID[@]}"; do
		flash_rom $index "$@"
		exitcode=$?
		[[ $exitcode -ne 0 ]] && ret=$exitcode
	done
	return $ret
}


print_usage() {
	echo "${CYAN}Usage:  nvidia-flash [command] [option] <vbios.rom>${NOCOLOR}

Commands:${GREEN}
${GREEN}  -l | --list               ${NOCOLOR}List supported GPU
${GREEN}  -w | --write ${BGREEN}INDEX/BUSID  ${NOCOLOR}Write vbios image to GPU
${GREEN}  --all                     ${NOCOLOR}Write vbios image to ${WHITE}all${NOCOLOR} NVIDIA gpu
${GREEN}  -s | --save ${BGREEN}INDEX/BUSID   ${NOCOLOR}Save vbios from GPU to file
${GREEN}  -b | --backup             ${NOCOLOR}Save vbioses from all GPU to files
${GREEN}  -v | --version            ${NOCOLOR}Print flasher version

Options:
${GREEN}  --force                   ${NOCOLOR}Force writing (use with --write or --all)

Examples:
${GREEN}  nvidia-flash -w 00:01 vbios.rom  ${NOCOLOR}Write bios.rom to GPU 00:01
${GREEN}  nvidia-flash --all --force vbios.rom  ${NOCOLOR}Force write bios.rom to all GPU
${GREEN}  nvidia-flash -s 0 vbios.rom      ${NOCOLOR}Save bios from GPU 0 to vbios.rom
${GREEN}  nvidia-flash --backup vbios.rom  ${NOCOLOR}Save bioses from all GPU to files vbios0.rom, vbios1.rom, ...

Warning:${YELLOW}
  NVIDIA driver and Xserver must be unloaded!
  Use ${WHITE}nvstop${YELLOW} for unload. ${WHITE}nvstop start${YELLOW} to resume.${NOCOLOR}"
}


# check CMD line for --force and remove it from the list
#CMD=()
#for(( i=0; i < ${#@}; i++ )); do
#	[[ $force -eq 0 && "${!i}" == "--force" ]] && echo "found at $i" && force=1 && continue
#	CMD+=( "${!i}" )
#done
#CMD+=( "${!i}" ) # do not check last element


case "$1" in
	-v|--version)
		nvflash_linux --list | grep -m1 -oP "\K[0-9\.]+"
	;;
	
	-l|--list)
		nvflash_linux --list | cat -s
	;;

	-w|--write)
		get_busid
		flash_rom "${@:2}"
	;;

	-s|--save)
		get_busid
		save_rom "${@:2}"
	;;

	--all)
		get_busid
		flash_all "${@:2}"
	;;

	-b|--backup)
		get_busid
		save_all "${@:2}"
	;;

	--busid)
		get_busid
		print_busid
	;;

	-h|--help|"")
		print_usage
	;;
esac

exit
