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

EAPI=8

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

MY_PN=torchaudio
MYPN=audio
MYP=${MYPN}-${PV}

DESCRIPTION="Audio processing for PyTorch (ROCm 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 +sox"
RESTRICT="test"

REQUIRED_USE="${PYTHON_REQUIRED_USE}"

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

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

DEPEND="${RDEPEND}"

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_ROCM=1
	export USE_CUDA=0
	export BUILD_SOX=$(usex sox 1 0)
	export USE_FFMPEG=$(usex ffmpeg 1 0)
	export BUILD_VERSION="${PV}"

	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

	# Fix RPATH on native extensions
	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
}