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

EAPI=8

DISTUTILS_USE_PEP517=setuptools
PYTHON_COMPAT=( python3_{13..14} )
LLVM_COMPAT=( 22 23 )

inherit distutils-r1 llvm-r2

DESCRIPTION="Tensor library and neural network framework with automatic differentiation"
HOMEPAGE="https://tinygrad.org/ https://github.com/tinygrad/tinygrad"
SRC_URI="https://github.com/tinygrad/tinygrad/archive/refs/tags/v${PV}.tar.gz -> ${P}.gh.tar.gz"

# The visualization assets also bundle highlight.js (BSD) and d3 (ISC).
LICENSE="BSD ISC MIT"
SLOT="0"
KEYWORDS="~amd64"
IUSE="cuda numpy opencl"

# The CPU backend invokes Clang at runtime, including for its worker kernels.
RDEPEND="
	$(llvm_gen_dep 'llvm-core/clang:${LLVM_SLOT}')
	cuda? ( dev-util/nvidia-cuda-toolkit )
	numpy? ( dev-python/numpy[${PYTHON_USEDEP}] )
	opencl? ( virtual/opencl )
"
BDEPEND="
	>=dev-python/setuptools-77[${PYTHON_USEDEP}]
"

EPYTEST_PLUGINS=()
EPYTEST_XDIST=1
distutils_enable_tests pytest

python_prepare_all() {
	distutils-r1_python_prepare_all

	# Slotted Clang is not necessarily available as /usr/bin/clang.
	# Keep the upstream CC environment override for callers that need it.
	sed -i \
		-e "s@getenv(\"CC\", 'clang')@getenv(\"CC\", '${EPREFIX}/usr/lib/llvm/${LLVM_SLOT}/bin/clang')@" \
		tinygrad/runtime/support/compiler_cpu.py || die
}

python_test() {
	local -x DEV=CPU
	local -x CC="$(get_llvm_prefix -b)/bin/clang"
	local -x XDG_CACHE_HOME="${T}/cache-${EPYTHON}"

	# Self-contained tensor, autograd, JIT, and dtype tests. The full suite
	# additionally needs GPUs, model downloads, and other ML frameworks.
	epytest test/test_tiny.py test/null/test_dtype.py
}

pkg_postinst() {
	elog "Use DEV=CPU to select the CPU backend explicitly."
	elog "GPU backends need compatible hardware and runtime libraries."
	elog "See https://docs.tinygrad.org/runtime/ for backend configuration."
}