#!/bin/ksh93

PKG_BASE_URL='https://developer.download.nvidia.com/compute/cuda/repos'
# TensorFlow
PKG_BASE_URL_TR='https://developer.download.nvidia.com/compute/machine-learning/repos'

typeset -r VERSION='1.0' FPROG=${.sh.file} PROG=${FPROG##*/} SDIR=${FPROG%/*}

function showUsage {
	[[ -n $1 ]] && X='-?' ||  X='--man'
	getopts -a ${PROG} "${ print ${USAGE} ; }" OPT $X
}

# takes a version string formatted as N(\.N)* from the first 3 numbers (N)
# (trailing ones are ignored) and returns a corresponding int value. N is
# expected to be in the range {0..999}. Can be used to get rid off damn stupid
# leading zeros and missing trailing zero parts.
function version2int {
	typeset -n VAL=$1
	typeset A=( ${2//./ } 0 0 0 )
	VAL=${A[2]}
	(( VAL+=${A[1]} * 1000 ))
	(( VAL+=$A * 1000000 ))
}

function getDownloadURL {
	typeset -a A
	[[ -z ${REL} ]] && A=( ${ lsb_release  -sir ; } ) || A=( Ubuntu ${REL} )
	typeset -l UV="${A}${A[1]//.}"
	[[ -n $1 ]] && print "${PKG_BASE_URL_TR}/${UV}/${ uname -m; }/" || \
	print "${PKG_BASE_URL}/${UV}/${ uname -m; }/"
}

function getAvailablePackages {
	typeset -n PKGS=$1
	typeset D=${ mktemp -d /tmp/nv.XXXXXX ; } X P F D
	[[ -z $D ]] && print -u2 'Exiting.' && return 1
	trap "rm -rf $D" EXIT

	for X in xsltproc wget ; do
		F=${ whence $X ; }
		[[ -z $F ]] && P+=", $X"
	done
	[[ -n $P ]] && print -u2 'Please install the following package(s) first:' \
		"${P:2} !" && return 2

	if [[ -n ${IDX} ]]; then
		cp ${IDX} $D/index.html
	else
		(( NO_PROXY )) && A='--no-proxy' || A=
		F="${ (( TR )) && getDownloadURL 1 || getDownloadURL; }/index.html"
		wget $A -qL -T 10 -O $D/index.html $F
		A=$?
		(( $A )) && print -u2 "wget could not get '$F' (exit code $A)." \
			'Try to fix this problem first.' && return 3

	fi
	# convert the output of an apache httpd auto index html file to a plain/text
    # list of files
	cat >$D/idx2file.xslt<<EOF
<?xml version="1.0" encoding="UTF-8"?>
<xsl:stylesheet version="1.0" xmlns:xsl="http://www.w3.org/1999/XSL/Transform">
<xsl:output method="text" encoding="UTF-8"/>
<xsl:strip-space elements="*"/>
<xsl:template match="@*|node()">
	<xsl:apply-templates select="node()" />
</xsl:template>
<xsl:template match='span[@class="file"]'>
	<xsl:value-of select='.'/><xsl:text>
</xsl:text>
</xsl:template>
</xsl:stylesheet>
EOF
	sed -i -e '0,/doctype html/ d' $D/index.html
	if (( LIST )); then
		xsltproc $D/idx2file.xslt $D/index.html
		return $?
	elif (( KMLIST )); then
		F=
		integer V
		xsltproc $D/idx2file.xslt $D/index.html | while read X ; do
			# irgendwas wie  nvidia-dkms-418_418.87.00-0ubuntu1_amd64.deb
			[[ ${X:0:12} == 'nvidia-dkms-' ]] || continue
			A=${X#*_}			# strip off package name and separator
			P=${A%%-*}			# strip off trailing ubuntu non-sense
			version2int V "$P"	# make it a single number
			F+="${X}:$V\n"		# map file to the single number
		done
		print -n "$F" | sort -t: -k2n,2n |cut -f1 -d:
		return 0
	fi
	X=
	xsltproc $D/idx2file.xslt $D/index.html | while read P ; do
		[[ ${P: -4:4} == '.deb' ]] && X+=" $P"
	done
	[[ -z $X ]] && print -u2 "The received contents for '$F' is unexpected." \
		"Try to fix this problem first." && return 4
	PKGS=( $X )
}

function getUniquePkgs {
	typeset -n PKGS=$1 PNAMES=$2 MAIN_VERS=$3 SUB_VERS=$4 VSEL=$5
	typeset -A N MV SV NSEL
	typeset -a S
	typeset A B C X P
	integer VERS V
	version2int VERS "$6"
	for X in ${PKGS[@]} ; do
		S=( ${X//_/ } )			# e.g. split cuda-10-0_10.0.130-1_amd64.deb
		A=${S%%*(-+([0-9]))}	# e.g. cuda-10-0 -> cuda
		B=${S:${#A}+1}			# -> 10-0
		N["$A"]=1
		MV[_"$B"]=1
		C="${S[1]%-*}"			# e.g. 10.0.130-1
		SV[_"$C"]=1
		version2int V "$C"
		(( V == VERS )) && VSEL["$A"]="${S}_${S[1]}:$V"
	done
	set -s -- ${!N[@]}
	PNAMES=( "$@" )
	set -s -- ${!MV[@]}
	MAIN_VERS=( "$@" )
	set -s -- ${!SV[@]}
	SUB_VERS=( "$@" )
	(( VERB )) || return 0
	X=
	for P in ${PNAMES[@]} ; do
		X+="$P\n"
	done
	print -u2 -n "\n${#PNAMES[@]} Unique package names:\n$X"
	X=
	for P in ${MAIN_VERS[@]} ; do
		X+="${P:1}\n"
	done
	print -u2 -n "\n${#MAIN_VERS[@]} Unique main versions:\n$X"
	X=
	for P in ${SUB_VERS[@]} ; do
		X+="${P:1}\n"
	done
	print -u2 -n "\n${#SUB_VERS[@]} Unique sub versions:\n$X"
	X=
	for P in ${!VSEL[@]} ; do
		X+="${VSEL[$P]%:*}\n"
	done
	print -u2 -n "\n${#VSEL[@]} Matching packages:\n$X"
	integer I
	X=
	for P in ${PNAMES[@]} ; do
		[[ -z ${VSEL[$P]} ]] && X+="$P\n" && (( I++ ))
	done
	print -u2 -n "\n$I Non-matching packages:\n$X"
}

# $1 .. vname: contains all packages from nvidia repo which are $2 compatible
# $2 .. the version string of the nvidia kernel module as is
function checkInstalled {
	typeset -n SEL=$1
	if (( ${#SEL[@]} == 0 )); then
		print 'The repository contains no packages for this version - exiting.'
		return 0
	fi

	integer FOUND LV RV
	typeset REV="$2" NREV KREV TBD PV SV KV OV CMDA CMDB MISSING A B V X Z
	A=( ${REV//./ } )
	NREV=${A##*(0)}
	NREV+=".${A[1]##*(0)}"
	X=${A[2]##*(0)}
	# normalized, i.e. not zero-padded version
	[[ -n $X ]] && NREV+=".$X" || KREV="0."

	V=${ uname -r ; }
	X=${V%-*}
	KREV+=${X//+([-])/.}			# space separated kernel version

	dpkg-query -W -f '${binary:Package}\t${Version}\n' | while read N OV ; do
		N=${N%:*}					# skip arch
		A=${N%%*(-+([0-9]))}		# skip "-XXX*" -> real package basename
		FOUND=0
		for S in ${!SEL[@]} ; do
			[[ $A == $S ]] && FOUND=1 && break
		done
		(( ! FOUND )) && continue

		V=${OV#*:}			# skip epoche overhead
		V=${V%%-*}			# skip trailing nonsense -> installed version

		# If the version of the installed package does not match ${REV} of the
		# corresponding package in the nvidia repo, one needs to install the
		# repo version of the package:
		SV=${SEL[$S]%:*}

		if [[ $A =~ ^(cuda-drivers|nvidia-(dkms|driver|headless-no-dkms))$ ]]
		then
			# But our modules are not zero padded, i.e. properly versioned, so
			# need to check for related versions in local apt cache
			PV=${NREV}
			[[ ${N:0:12} == 'nvidia-dkms-' ]] && PV+=".${KREV}"

			B=
			LV=-1
			apt-cache madison "${N}" | while read -A A ; do
				X=${A[2]%%-*}				# strip off pkg-release|debian crap
				Z=${.sh.match:1}			# release info
				# ignore foreign packages, even if correct version
				[[ $X == ${PV} && $Z =~ ^[0-9]+$ ]] || continue

				# easy to determine the latest package release
				(( $Z > LV )) && LV=$Z && B="${A[2]}"
			done
			[[ ${OV} == $B ]] && continue
			[[ -z $B ]] && MISSING+="\\\\\n\t${SV//_*}-${PV} " && B="${PV}-1"
			if [[ ${N:0:12} == 'nvidia-dkms-' ]]; then
				CMDA+="\tapt-get install ${SV//_*}=$B\n"
			else
				CMDB+=" \\\\\n\t\t${SV//_*}=$B"
			fi
		else
			# installed package version == repo package version ? :
			[[ "${N}_$V" == ${SV%-*} ]] && continue
			# nope
			CMDB+=" \\\\\n\t\t${SV//_/=}"
		fi
		TBD+="\t${N}: ${OV}\n"
	done
	if [[ -z ${TBD} ]]; then
		print -u2 'No packages with version mismatch found.'
	else
		print '\nThe following installed packages are not compatible with' \
			"version ${REV}:\n${TBD}"
		[[ -n ${MISSING} ]] && print "Missing the following packages in" \
			"your apt cache (forgot 'apt-get update'?):${MISSING:2}\n"
		[[ -n ${CMDA} && -n ${CMDB} ]] && X='commands' || X='command'
		print "The following $X may help to install a compatible package set:\n"
		[[ -n ${CMDA} ]] && print "${CMDA}"
		[[ -n ${CMDB} ]] && print "\tapt-get install ${CMDB}"
	fi
}

function doMain {
	if (( SURL )); then
		(( TR )) && getDownloadURL 1 || getDownloadURL
		return $?
	fi
	typeset F='/proc/driver/nvidia/version'
	typeset -a NVPKGS PN MV SV
	typeset -A SEL
	if [[ -z ${REV} ]]; then
		if [[ ! -e $F ]]; then
			print -u2 'There is no nvidia kernel module loaded,' \
				'so nothing to check.'
			return 1
		fi
		typeset X=$(<$F)
		X=${X##*Kernel Module*( )}
		X=${X%% *}
		if [[ ! $X =~ ^[0-9]+(\.[0-9]+)*$ ]]; then
			print -u2 'Unable to determine the version of the currently' \
				'loaded nvidia module.'
			return 2
		fi
		REV="$X"
	fi
	(( ! (LIST | KMLIST) )) && print -u2 "Checking for version ${REV} ..."
	getAvailablePackages NVPKGS || return 3
	(( LIST )) && return 0
	getUniquePkgs NVPKGS PN MV SV SEL ${REV} || return 3
	checkInstalled SEL ${REV}
}

USAGE="[-?${VERSION}"' ]
[-copyright?Copyright (c) 2019 Jens Elkner. All rights reserved.]
[-license?CDDL 1.0]
[+NAME?'"${PROG}"' - check nvidia kernel module version vs. installed packages.]
[+DESCRIPTION?Determines the version of the currently installed nvidia kernel module and checks it against installed nvidia packages for ubuntu. On mismatch a command gets generated and shown, which should fix the mismatch on execution.]
[+NOTES?Make sure, that the script is able to access '"${PKG_BASE_URL}"/' and '"${PKG_BASE_URL_TR}"/'.]
[h:help?Print this help and exit.]
[F:functions?Print a list of all functions available.]
[T:trace]:[functionList?A comma separated list of functions of this script to trace (convinience for troubleshooting).]
[+?]
[L:list?Just show the list of all downloadable files and exit.]
[R:release]:[release?Use the given Ubuntu release to determine repo URLs and packages instead of `lsb_release -sr`. E.g. "18.04".]
[i:index]:[file?Use the given file instead of downloading the current index.html of the related package repository.]
[k:kmods?Show all kernel module packages available from nvidia and exit. Might be used to determine, how far away the installed version is.]
[p:proxy?Do not pass the \b--no-proxy\b option to wget, when used.]
[r:revision]:[version?Use the given nvidia kernel module \aversion\a for comparision instead of using the one of the currently loaded nvidia kernel module (e.g. 418.87.01).]
[t:tensor?Apply the given arguments having tensor instead of cuda runtime in mind.]
[u:url?Show the URL to use for downloading original nvidia packages for ubuntu and exit.]
[v:verbose?Print more annoying details about the things the script is doing.]
'

X="${ print ${USAGE} ; }"
integer NO_PROXY=1 VERB=0 LIST=0 KMLIST=0 SURL=0 TR=0
unset IDX REV REL; IDX= REV= REL=
while getopts "${X}" OPT ; do
	case ${OPT} in
		h) showUsage ; exit 0 ;;
		T)	if [[ ${OPTARG} == 'ALL' ]]; then
				typeset -ft ${ typeset +f ; }
			else
				typeset -ft ${OPTARG//,/ }
			fi
			;;
		F) typeset +f && exit 0 ;;
		L) LIST=1 ;;
		R) REL="${OPTARG}" ;;
		i) IDX="${OPTARG}" ;;
		k) KMLIST=1 ;;
		p) NO_PROXY=0 ;;
		r) REV=${OPTARG} ;;
		v) VERB=1 ;;
		u) SURL=1 ;;
		t) TR=1 ;;
		*) showUsage 1 ; exit 1 ;;
	esac
done

X=$((OPTIND-1))
shift $X && OPTIND=1
unset X

doMain "$@"