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

EAPI=8

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

FUNCHOOK_PV=1.1.3

inherit distutils-r1

DESCRIPTION="AI Model Dynamic Offloader for ComfyUI"
HOMEPAGE="https://github.com/Comfy-Org/comfy-aimdo"
SRC_URI="
	https://github.com/Comfy-Org/${PN}/archive/refs/tags/v${PV}.tar.gz -> ${P}.gh.tar.gz
	https://github.com/kubo/funchook/releases/download/v${FUNCHOOK_PV}/funchook-${FUNCHOOK_PV}.tar.gz
"

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

LICENSE="GPL-3"
SLOT="0"
KEYWORDS="~amd64"
IUSE="+cuda +rocm"

RDEPEND="
	|| (
		sci-ml/pytorch[${PYTHON_SINGLE_USEDEP}]
		sci-ml/pytorch-rocm[${PYTHON_SINGLE_USEDEP}]
		sci-ml/pytorch-cuda[${PYTHON_SINGLE_USEDEP}]
	)
"
BDEPEND="
	dev-build/cmake
"

export SETUPTOOLS_SCM_PRETEND_VERSION=${PV}

src_compile() {
	local funchook_src="${WORKDIR}/funchook-${FUNCHOOK_PV}"
	local funchook_build="${T}/funchook-build"

	# Build funchook static lib
	cmake -S "${funchook_src}" -B "${funchook_build}" \
		-DCMAKE_BUILD_TYPE=Release \
		-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
		-DFUNCHOOK_BUILD_SHARED=OFF \
		-DFUNCHOOK_BUILD_STATIC=ON \
		-DFUNCHOOK_BUILD_TESTS=OFF \
		-DFUNCHOOK_DISASM=distorm \
		-DFUNCHOOK_INSTALL=OFF \
		|| die "funchook cmake failed"
	cmake --build "${funchook_build}" -j$(makeopts_jobs) || die "funchook build failed"

	local funchook_libs="${funchook_build}/libfunchook.a ${funchook_build}/libdistorm.a"

	if use cuda; then
		gcc -shared -o "${S}/comfy_aimdo/aimdo.so" -fPIC -O2 -g -pthread \
			"${S}"/src/*.c "${S}"/src-cuda/dispatch.c "${S}"/src-posix/*.c \
			-I"${S}/src" -I"${funchook_src}/include" \
			${funchook_libs} -ldl || die "CUDA .so build failed"
	fi

	if use rocm; then
		gcc -shared -o "${S}/comfy_aimdo/aimdo_rocm.so" -fPIC -O2 -g -pthread \
			-D__HIP_PLATFORM_AMD__ \
			"${S}"/src/*.c "${S}"/src-hip/dispatch.c "${S}"/src-posix/*.c \
			-I"${S}/src" -I"${funchook_src}/include" \
			${funchook_libs} -ldl || die "ROCm .so build failed"
	fi

	# Write version file for setuptools-scm
	cat > "${S}/comfy_aimdo/_version.py" <<-PYEOF
	version = "${PV}"
	__version__ = version
	__version_tuple__ = tuple(version.split("."))
	PYEOF

	distutils-r1_src_compile
}