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

EAPI=8

DISTUTILS_USE_PEP517=setuptools
DISTUTILS_SINGLE_IMPL=1
DISTUTILS_EXT=1
PYTHON_COMPAT=( python3_{12..14} )

inherit cuda distutils-r1

DESCRIPTION="Decode/encode video and audio into PyTorch tensors via FFmpeg"
HOMEPAGE="
	https://github.com/meta-pytorch/torchcodec
	https://pypi.org/project/torchcodec/
"
SRC_URI="
	https://github.com/meta-pytorch/torchcodec/archive/refs/tags/v${PV}.tar.gz
		-> ${P}.gh.tar.gz
"

LICENSE="BSD-2"
SLOT="0"
KEYWORDS="~amd64 ~arm64"
IUSE="cuda"

# Upstream dlopens FFmpeg 4-8; ::gentoo supplies the live 7/8 slot.
RDEPEND="
	sci-ml/pytorch[${PYTHON_SINGLE_USEDEP}]
	media-video/ffmpeg:=
	cuda? (
		dev-util/nvidia-cuda-toolkit:=
	)
"
DEPEND="${RDEPEND}"
BDEPEND="
	$(python_gen_cond_dep '
		dev-python/pybind11[${PYTHON_USEDEP}]
	')
	dev-build/cmake
"

# Tests pull a video corpus from S3.
RESTRICT="test"

python_compile() {
	# setup.py drives CMake through environment variables.
	export CMAKE_BUILD_TYPE=Release
	export BUILD_VERSION="${PV}"

	# Link system FFmpeg instead of the S3-vendored wheel payload; acknowledge
	# upstream's redistribution guard for this local source build.
	export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1

	# Torch's CMake enables CUDA whenever the toolkit is installed, even for CPU
	# builds, and nvcc rejects a gcc newer than its toolkit supports. Let
	# cuda_gccdir pick the newest installed gcc it accepts. verified 2026-10-08
	if use cuda || has_version -b dev-util/nvidia-cuda-toolkit; then
		# It returns nothing under a non-GNU toolchain, and a path that does
		# not exist where CHOST is not *-pc-linux-gnu; nvcc then keeps its
		# default host compiler.
		local cuda_gcc_bindir
		if cuda_gcc_bindir="$(cuda_gccdir)" && [[ -x ${cuda_gcc_bindir}/g++ ]]; then
			export CUDAHOSTCXX="${cuda_gcc_bindir}/g++"
		fi
	fi

	if use cuda; then
		export ENABLE_CUDA=1

		# With no list, Torch's CMake emits gencode flags for every architecture
		# it knows, down to compute_50, which CUDA 13 dropped. Carry PTX with the
		# fallback: a cubin alone would not load on a GPU newer than the default.
		# verified 2026-09-16
		if [[ -z ${TORCH_CUDA_ARCH_LIST} ]]; then
			ewarn "WARNING: torchcodec is being built for CUDA compute capability 7.5,"
			ewarn "plus PTX so the driver can JIT for a newer GPU. This default may"
			ewarn "not be optimal for your GPU."
			ewarn ""
			ewarn "To target your GPU, set TORCH_CUDA_ARCH_LIST through Portage's package.env"
			ewarn "mechanism and re-emerge torchcodec."
			ewarn ""
			ewarn "You can look up your GPU's CUDA compute capability at https://developer.nvidia.com/cuda-gpus"
			ewarn "or by running /opt/cuda/extras/demo_suite/deviceQuery | grep 'CUDA Capability'"
		fi
		export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-7.5+PTX}"
	else
		export ENABLE_CUDA=
	fi

	distutils-r1_python_compile
}