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

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

LIBDIR="${EROOT}/usr/lib/caffe2"
LIB64="${EROOT}/usr/lib64"
INCDIR="${EROOT}/usr/include"
CMAKEDIR="${EROOT}/usr/share/cmake"
BINDIR="${EROOT}/usr/bin"

# Libraries that get symlinked
CAFFE2_LIBS=(
	libc10.so
	libtorch.so
	libtorch_cpu.so
	libtorch_python.so
	libtorch_global_deps.so
	libshm.so
	libcaffe2_nvrtc.so
)

# Backend-specific libs
ROCM_LIBS=( libc10_hip.so libtorch_hip.so )
CUDA_LIBS=( libc10_cuda.so libtorch_cuda.so )

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

get_active() {
	if [[ -L "${LIB64}/libtorch.so" ]]; then
		local target
		target=$(readlink "${LIB64}/libtorch.so")
		target=${target#../lib/caffe2/}
		target=${target%%/*}
		echo "${target}"
	fi
}

remove_symlinks() {
	local lib
	for lib in "${CAFFE2_LIBS[@]}" "${ROCM_LIBS[@]}" "${CUDA_LIBS[@]}"; do
		rm -f "${LIB64}/${lib}"
	done
	rm -f "${BINDIR}/torch_shm_manager"

	# Include dirs
	local inc
	for inc in ATen c10 caffe2 torch; do
		rm -f "${INCDIR}/${inc}"
	done

	# CMake dirs
	local cm
	for cm in ATen Caffe2 Torch; do
		rm -f "${CMAKEDIR}/${cm}"
	done
}

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

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

	# Libraries
	local lib
	for lib in "${prefix}"/lib64/*.so; do
		[[ -f "${lib}" ]] || continue
		ln -s "../lib/caffe2/${target}/lib64/$(basename "${lib}")" "${LIB64}/$(basename "${lib}")" 2>/dev/null
	done

	# Binary
	[[ -f "${prefix}/bin/torch_shm_manager" ]] && \
		ln -s "../lib/caffe2/${target}/bin/torch_shm_manager" "${BINDIR}/torch_shm_manager"

	# Includes
	local inc
	for inc in "${prefix}"/include/*/; do
		[[ -d "${inc}" ]] || continue
		ln -s "../lib/caffe2/${target}/include/$(basename "${inc}")" "${INCDIR}/$(basename "${inc}")" 2>/dev/null
	done

	# CMake
	local cm
	for cm in "${prefix}"/share/cmake/*/; do
		[[ -d "${cm}" ]] || continue
		ln -s "../../lib/caffe2/${target}/share/cmake/$(basename "${cm}")" "${CMAKEDIR}/$(basename "${cm}")" 2>/dev/null
	done
}

### list action ###

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

do_list() {
	local active
	active=$(get_active)
	write_list_start "Available caffe2 backends:"

	local d i=1
	for d in "${LIBDIR}"/*/; do
		[[ -d "${d}lib64" ]] || continue
		local variant="$(basename "${d}")"
		if [[ "${variant}" == "${active}" ]]; then
			write_kv_list_entry "[${i}]" "${variant} *"
		else
			write_kv_list_entry "[${i}]" "${variant}"
		fi
		((i++))
	done
}

### show action ###

describe_show() {
	echo "Show currently active caffe2 backend"
}

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

### set action ###

describe_set() {
	echo "Set active caffe2 backend"
}

describe_set_parameters() {
	echo "<target>"
}

do_set() {
	[[ -z "${1}" ]] && die -q "You must specify a backend (rocm, cuda)"

	local target="${1}"

	if is_number "${target}"; then
		local targets=( $(find_targets) )
		target="${targets[$((${1} - 1))]}"
		[[ -z "${target}" ]] && die -q "Invalid selection"
	fi

	remove_symlinks
	set_symlinks "${target}"
}