# Copyright 2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

DESCRIPTION="Manage torchvision backend (ROCm/CUDA)"
VERSION="0.1"

TORCHVISIONDIR="${EROOT}/usr/lib/torchvision"
SITEDIR="${EROOT}/usr/lib/python3.14/site-packages"

find_targets() {
	local targets=()
	local d
	for d in "${TORCHVISIONDIR}"/*/; do
		[[ -d "${d}torchvision" ]] && targets+=("$(basename "${d}")")
	done
	echo "${targets[@]}"
}

get_active() {
	if [[ -L "${SITEDIR}/torchvision" ]]; then
		local target
		target=$(readlink "${SITEDIR}/torchvision")
		target=${target#../../torchvision/}
		target=${target%%/*}
		echo "${target}"
	fi
}

remove_symlinks() {
	rm -f "${SITEDIR}/torchvision"
}

set_symlinks() {
	local target="${1}"
	local prefix="${TORCHVISIONDIR}/${target}"

	if [[ ! -d "${prefix}/torchvision" ]]; then
		die -q "Target '${target}' not found: ${prefix}/torchvision"
	fi

	ln -s "../../torchvision/${target}/torchvision" "${SITEDIR}/torchvision"
}

### list action ###

describe_list() {
	echo "List available torchvision backends"
}

do_list() {
	local targets=( $(find_targets) )
	local active=$(get_active)
	local i

	write_list_start "Available torchvision backends:"
	for (( i = 0; i < ${#targets[@]}; i++ )); do
		if [[ "${targets[i]}" == "${active}" ]]; then
			targets[i]=$(highlight_marker "${targets[i]}")
		fi
	done
	write_numbered_list -m "(none available)" "${targets[@]}"
}

### show action ###

describe_show() {
	echo "Show the current torchvision backend"
}

do_show() {
	local active=$(get_active)
	echo "${active:-(unset)}"
}

### set action ###

describe_set() {
	echo "Set the torchvision backend"
}

describe_set_parameters() {
	echo "<target>"
}

describe_set_options() {
	echo "target : Target name or number (from 'list' action)"
}

do_set() {
	local target="${1}"

	if [[ -z "${target}" ]]; then
		die -q "Required parameter 'target' missing"
	fi

	local targets=( $(find_targets) )
	if is_number "${target}"; then
		target=${targets[$((target - 1))]}
	fi

	local found=0
	local t
	for t in "${targets[@]}"; do
		[[ "${t}" == "${target}" ]] && found=1
	done
	if [[ ${found} -eq 0 ]]; then
		die -q "Invalid target '${target}'"
	fi

	remove_symlinks
	set_symlinks "${target}"
}