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

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=torchaudio
MYPN=audio
MYP=${MYPN}-${PV}

DESCRIPTION="Audio processing for PyTorch (CUDA backend)"
HOMEPAGE="https://github.com/pytorch/audio"
SRC_URI="https://github.com/pytorch/${MYPN}/archive/refs/tags/v${PV}.tar.gz
	-> torchaudio-${PV}.tar.gz"

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

LICENSE="BSD-2"
SLOT="0"
KEYWORDS="~amd64"
IUSE="+ffmpeg"
RESTRICT="test"

REQUIRED_USE="${PYTHON_REQUIRED_USE}"

CAFFE2_PREFIX="/usr/lib/caffe2/cuda"
PYTORCH_PREFIX="/usr/lib/pytorch/cuda"
TORCHAUDIO_PREFIX="/usr/lib/torchaudio/cuda"

RDEPEND="
	${PYTHON_DEPS}
	$(python_gen_cond_dep '
		dev-python/numpy[${PYTHON_USEDEP}]
	')
	ffmpeg? ( media-video/ffmpeg )
	~sci-ml/caffe2-cuda-2.12.0[${PYTHON_SINGLE_USEDEP}]
	~sci-ml/pytorch-cuda-2.12.0[${PYTHON_SINGLE_USEDEP}]
"

DEPEND="${RDEPEND}"

src_prepare() {
	cuda_src_prepare
	distutils-r1_src_prepare
}

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

	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	export USE_CUDA=1
	export USE_ROCM=0
	export USE_FFMPEG=$(usex ffmpeg 1 0)
	export BUILD_VERSION="${PV}"
	export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-6.1 7.5}"

	cuda_add_sandbox
	addpredict "/dev/char/"

	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} --expt-relaxed-constexpr"

	# CRITICAL: Point cmake find_package(Torch) to CUDA config, not eselect ROCm
	export Torch_DIR="${CAFFE2_PREFIX}/share/cmake/Torch"
	export Caffe2_DIR="${CAFFE2_PREFIX}/share/cmake/Caffe2"
	export CMAKE_PREFIX_PATH="${CAFFE2_PREFIX}:${CMAKE_PREFIX_PATH}"

	export CPLUS_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${CPLUS_INCLUDE_PATH}"
	export C_INCLUDE_PATH="${CAFFE2_PREFIX}/include:${C_INCLUDE_PATH}"
	export CUDAFLAGS="-I${CAFFE2_PREFIX}/include --expt-relaxed-constexpr"

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

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

	export PYTHONPATH="${PYTORCH_PREFIX}:${PYTHONPATH}"
	export LD_LIBRARY_PATH="${CAFFE2_PREFIX}/lib64:${LD_LIBRARY_PATH}"

	distutils-r1_python_install

	local sitedir="$(python_get_sitedir)"
	if [[ -d "${ED}${sitedir}/torchaudio" ]]; then
		mv "${ED}${sitedir}/torchaudio" "${dest}/torchaudio" || die
	fi

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

	local cso
	for cso in "${dest}"/torchaudio/lib/*.so; do
		[[ -f "${cso}" ]] && patchelf --set-rpath "${CAFFE2_PREFIX}/lib64:\$ORIGIN" "${cso}" 2>/dev/null || true
	done
}