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

EAPI=8

ROCM_VERSION=7.2
LLAMA_CPP_COMMIT="8c146a8366304c871efc26057cc90370ccf58dad"

inherit cmake go-module rocm systemd

DESCRIPTION="Get up and running with large language models."
HOMEPAGE="https://ollama.com"
SRC_URI="
	https://github.com/ollama/${PN}/archive/refs/tags/v${PV}.tar.gz -> ${P}.gh.tar.gz
	https://vendors.simple-co.de/${PN}/${PN}-${PV}-deps.tar.xz
	https://github.com/ggml-org/llama.cpp/archive/${LLAMA_CPP_COMMIT}.tar.gz -> llama-cpp-${LLAMA_CPP_COMMIT:0:8}.tar.gz
"
S="${WORKDIR}/${PN}-${PV}"
LICENSE="MIT"
SLOT="0"

KEYWORDS="~amd64"

IUSE="cuda systemd video_cards_amdgpu cpu_flags_x86_avx cpu_flags_x86_avx2
cpu_flags_x86_avx512f cpu_flags_x86_avx512vbmi cpu_flags_x86_avx512_vnni
cpu_flags_x86_avx512_bf16
"

REQUIRED_USE="
	cpu_flags_x86_avx2? ( cpu_flags_x86_avx )
	cpu_flags_x86_avx512f? ( cpu_flags_x86_avx2 )
	cpu_flags_x86_avx512vbmi? ( cpu_flags_x86_avx512f )
	cpu_flags_x86_avx512_vnni? ( cpu_flags_x86_avx512f )
	cpu_flags_x86_avx512_bf16? ( cpu_flags_x86_avx512f )
"

RDEPEND="
	acct-group/ollama
	acct-user/ollama
"

DEPEND="
	>=dev-lang/go-1.25.5
	>=dev-build/cmake-3.31.9
	>=sys-devel/gcc-11.5.0
	cuda? (
		dev-util/nvidia-cuda-toolkit
		dev-libs/cccl
	)
	video_cards_amdgpu? ( sci-libs/hipBLAS[${ROCM_USEDEP}] )
	systemd? ( sys-apps/systemd )
"

pkg_pretend() {
	if use video_cards_amdgpu && [[ -d /opt/rocm ]]; then
		eerror "/opt/rocm exists and will interfere with HIP compiler detection."
		eerror "Gentoo ROCm lives under /usr — remove or unmount /opt/rocm first."
		die "/opt/rocm must not exist during ROCm builds"
	fi
}

src_prepare() {
	# Fix multilib path: upstream hardcodes "lib" but Gentoo amd64 uses "lib64".
	# Without this, the ollama binary cannot locate llama-server and GPU backends.
	local libdir="$(get_libdir)"
	if [[ "${libdir}" != "lib" ]]; then
		sed -i \
			-e 's|exeDir, "\.\.", "lib", "ollama"|exeDir, "..", "'"${libdir}"'", "ollama"|g' \
			-e 's|exeDir, "lib", "ollama"|exeDir, "'"${libdir}"'", "ollama"|g' \
			ml/path.go || die "failed to patch ml/path.go for multilib"
	fi

	# Apply ollama compat patches to the llama.cpp source tree.
	# When providing llama.cpp via FETCHCONTENT_SOURCE_DIR, the superbuild
	# passes OLLAMA_LLAMA_CPP_SKIP_COMPAT_PATCH=ON to sub-builds, so
	# patches must already be applied in the source.
	local llama_src="${WORKDIR}/llama.cpp-${LLAMA_CPP_COMMIT}"
	local p
	for p in $(find "${S}/llama/compat" -name "*.patch" -type f | sort); do
		einfo "Applying compat patch: $(basename ${p})"
		patch -d "${llama_src}" -p1 < "${p}" || die "compat patch failed: ${p}"
	done


	# Fix hipblasSgemm crash on RDNA4 (gfx1201) for F32 matrix-matrix multiply.
	# hipblasSgemm with HIPBLAS_OP_T fails for N>1 on gfx1201. Route through
	# hipblasGemmEx which uses different rocBLAS kernels.
	if use video_cards_amdgpu; then
		einfo "Applying hipblasSgemm RDNA4 workaround patch"
		patch -d "${llama_src}" -p1 < "${FILESDIR}/llama-cpp-hip-sgemm-gfx1201.patch" \
			|| die "hipblasSgemm patch failed"
	fi

	# Fix SOLVE_TRI crash on ROCm gfx1201 (and other GPUs without rocBLAS
	# strsm kernels). Extends the custom kernel to handle k>32 via tiling,
	# bypassing rocBLAS entirely. Affects Qwen3.5 and other DeltaNet models.
	# Upstream: https://github.com/ggml-org/llama.cpp/issues/19442
	if use video_cards_amdgpu; then
		einfo "Applying solve_tri ROCm tiling patch"
		patch -d "${llama_src}" -p1 < "${FILESDIR}/llama-cpp-solve-tri-rocm-tiling.patch" \
			|| die "solve_tri patch failed"
	fi

	cmake_src_prepare
}

src_configure() {
	# Ensure HIP detection uses system paths (/usr), not /opt/rocm
	export ROCM_PATH="${EPREFIX}/usr"
	export HIP_PATH="${EPREFIX}/usr"

	# Ollama's superbuild spawns sub-builds via "cmake --build" which do NOT
	# inherit ninja's -j flag. Without this, sub-builds run sequentially.
	# Cap at 16 to avoid swap thrashing during GPU compilation (hipcc/nvcc
	# use ~2GB per process; 16×2 + ~20GB baseline = ~52GB of 64GB).
	export CMAKE_BUILD_PARALLEL_LEVEL=16

	# Build backend list from USE flags
	local backends=()
	use cuda && backends+=( cuda_v12 )
	use video_cards_amdgpu && backends+=( rocm_v7_2 )
	local backend_list
	backend_list=$(IFS=";"; echo "${backends[*]}")

	local mycmakeargs=(
		-DOLLAMA_VERSION="${PV}"
		-DOLLAMA_LIB_DIR="$(get_libdir)/ollama"
		-DOLLAMA_MLX_BACKENDS=""
		-DOLLAMA_LLAMA_BACKENDS="${backend_list}"
		-DFETCHCONTENT_SOURCE_DIR_LLAMA_CPP="${WORKDIR}/llama.cpp-${LLAMA_CPP_COMMIT}"
		-DOLLAMA_LLAMA_CPP_SKIP_COMPAT_PATCH=ON
		-DOLLAMA_LLAMA_CPP_COMPAT_DIR="${S}/llama/compat"
		-DGGML_CCACHE=OFF
		-DCMAKE_SKIP_BUILD_RPATH=ON
		-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON
		-DCMAKE_INSTALL_RPATH="\$ORIGIN"
	)

	# AMD ROCm
	# Ollama sets "-parallel-jobs=4" internally but clang >=22 renamed it
	# to "--offload-jobs". Override with the correct flag for our toolchain.
	if use video_cards_amdgpu; then
		mycmakeargs+=(
			-DCMAKE_HIP_ARCHITECTURES="${AMDGPU_TARGETS// /;}"
			-DAMDGPU_TARGETS="${AMDGPU_TARGETS// /;}"
			-DCMAKE_PREFIX_PATH="${EPREFIX}/usr"
			-DCMAKE_HIP_FLAGS="--offload-jobs=4 -Wno-ignored-attributes -Wno-deprecated-pragma"
		)
	fi

	# nVidia CUDA
	if use cuda; then
		# Build only for targets set in CUDA_TARGETS (make.conf).
		# Fallback to sm_75 if nothing set.
		local cuda_arches="${CUDA_TARGETS:-sm_75}"
		# Strip "sm_" prefix and join with semicolons for cmake
		cuda_arches="${cuda_arches//sm_/}"
		cuda_arches="${cuda_arches// /;}"

		export CUDAHOSTCXX=/usr/bin/x86_64-pc-linux-gnu-gcc-14
		mycmakeargs+=(
			-DCMAKE_CUDA_COMPILER=/opt/cuda/bin/nvcc
			-DCMAKE_CUDA_HOST_COMPILER="${CUDAHOSTCXX}"
			-DCMAKE_CUDA_ARCHITECTURES="${cuda_arches}"
		)
	fi

	cmake_src_configure
}

src_install() {
	cmake_src_install

	dodir /etc/ld.so.conf.d
	echo "/usr/$(get_libdir)/ollama" > "${ED}/etc/ld.so.conf.d/ollama.conf"

	use systemd && systemd_dounit "${FILESDIR}/ollama.service"
}

pkg_preinst() {
	keepdir /var/log/ollama
	fowners ollama:ollama /var/log/ollama

	keepdir /var/lib/ollama/models
	fowners ollama:ollama /var/lib/ollama/models
}

pkg_postinst() {
	einfo "Quick guide:"
	einfo "systemctl start ollama"
	einfo "ollama run deepseek-r1:7b"
	einfo "See available models at https://ollama.com/library"
}