# Copyright 2024-2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2
# CUDA implementation of torchvision

EAPI=8

PYTHON_COMPAT=( python3_{11..14} )
DISTUTILS_SINGLE_IMPL=1
DISTUTILS_USE_PEP517=setuptools
DISTUTILS_EXT=1
inherit cuda distutils-r1 multiprocessing

MY_PN=torchvision
MYPN=vision
MYP=${MYPN}-${PV}

DESCRIPTION="Datasets, transforms and models for computer vision (CUDA backend)"
HOMEPAGE="https://github.com/pytorch/vision"
SRC_URI="https://github.com/pytorch/${MYPN}/archive/refs/tags/v${PV}.tar.gz
	-> torchvision-${PV}.tar.gz"

S="${WORKDIR}"/${MYP}

LICENSE="BSD"
SLOT="0"
KEYWORDS="~amd64"
IUSE="+ffmpeg +jpeg +png +webp"
RESTRICT="test"

REQUIRED_USE="${PYTHON_REQUIRED_USE}"

CAFFE2_PREFIX="/usr/lib/caffe2/cuda"
PYTORCH_PREFIX="/usr/lib/pytorch/cuda"
TORCHVISION_PREFIX="/usr/lib/torchvision/cuda"

RDEPEND="
	${PYTHON_DEPS}
	$(python_gen_cond_dep '
		dev-python/numpy[${PYTHON_USEDEP}]
		dev-python/pillow[${PYTHON_USEDEP}]
	')
	jpeg? ( media-libs/libjpeg-turbo:= )
	png? ( media-libs/libpng:= )
	webp? ( media-libs/libwebp )
	ffmpeg? ( media-video/ffmpeg )
	~sci-ml/caffe2-cuda-2.12.0[${PYTHON_SINGLE_USEDEP}]
	~sci-ml/pytorch-cuda-2.12.0[${PYTHON_SINGLE_USEDEP}]
	app-eselect/eselect-torchvision
"

DEPEND="${RDEPEND}"

src_prepare() {
	cuda_src_prepare
	distutils-r1_src_prepare
}

python_compile() {
	addpredict /dev/kfd
	addpredict /dev/random

	# Point torch imports to CUDA backend
	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	export FORCE_CUDA=1
	export TORCHVISION_USE_PNG=$(usex png 1 0)
	export TORCHVISION_USE_JPEG=$(usex jpeg 1 0)
	export TORCHVISION_USE_WEBP=$(usex webp 1 0)
	export TORCHVISION_USE_FFMPEG=$(usex ffmpeg 1 0)
	export TORCHVISION_USE_NVJPEG=1
	export TORCHVISION_USE_VIDEO_CODEC=1
	export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-6.1 7.5}"

	cuda_add_sandbox
	addpredict "/dev/char/"

	# CUDA 12.9 requires GCC < 15 — torch cpp_extension.py checks CXX directly
	local gcc14="/usr/x86_64-pc-linux-gnu/gcc-bin/14"
	export CC="${gcc14}/x86_64-pc-linux-gnu-gcc"
	export CXX="${gcc14}/x86_64-pc-linux-gnu-g++"
	export CUDAHOSTCXX="${CXX}"
	export NVCC_FLAGS="${NVCCFLAGS} -ccbin ${gcc14}"

	# CUDA caffe2 headers must be found before eselect ROCm symlinks in /usr/include
	export CPLUS_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${CPLUS_INCLUDE_PATH}"
	export C_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${C_INCLUDE_PATH}"

	MAX_JOBS="$(get_makeopts_jobs)" \
		distutils-r1_python_compile -j1
}

python_install() {
	local dest="${ED}${TORCHVISION_PREFIX}"
	mkdir -p "${dest}" || die

	# Need PYTHONPATH for install step too (imports torch)
	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	distutils-r1_python_install

	# Move from site-packages to our prefix
	local sitedir="$(python_get_sitedir)"
	if [[ -d "${ED}${sitedir}/torchvision" ]]; then
		mv "${ED}${sitedir}/torchvision" "${dest}/torchvision" || die
	fi

	# Move dist-info
	local distinfo
	for distinfo in "${ED}${sitedir}"/torchvision-*.dist-info; do
		[[ -d "${distinfo}" ]] && mv "${distinfo}" "${dest}/" || true
	done

	# Fix RPATH on _C.so to find CUDA caffe2 libs
	local cso="${dest}/torchvision/_C.so"
	if [[ -f "${cso}" ]]; then
		patchelf --set-rpath "${CAFFE2_PREFIX}/lib64:\$ORIGIN" "${cso}" || \
			ewarn "patchelf failed on _C.so — RPATH may need manual fix"
	fi
}

pkg_postinst() {
	local active
	active=$(eselect torchvision show 2>/dev/null)
	if [[ "${active}" == "(unset)" || -z "${active}" ]]; then
		eselect torchvision set cuda
		elog "torchvision backend set to: cuda"
	fi
}