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

EAPI=8

ROCM_VERSION=6.1
inherit git-r3 go-module rocm


DESCRIPTION="Get up and running with Llama 3, Mistral, Gemma, and other language models."
HOMEPAGE="https://ollama.com"
EGIT_REPO_URI="https://github.com/ollama/ollama.git"
LICENSE="MIT"
SLOT="0"

IUSE="cuda video_cards_amdgpu
cpu_flags_x86_avx cpu_flags_x86_avx2 cpu_flags_x86_avx512
"

REQUIRED_USE="
	cpu_flags_x86_avx2? ( cpu_flags_x86_avx )
	cpu_flags_x86_avx512? ( cpu_flags_x86_avx2 )
"

RDEPEND="
	acct-group/ollama
	acct-user/ollama
"
IDEPEND="${RDEPEND}"
BDEPEND="
	>=dev-lang/go-1.22.0
	>=dev-build/cmake-3.24
	>=sys-devel/gcc-11.4.0
	cuda? ( dev-util/nvidia-cuda-toolkit )
	video_cards_amdgpu? (
		=sci-libs/hipBLAS-${ROCM_VERSION}*
	)
"

pkg_pretend() {
	if use video_cards_amdgpu; then
		ewarn "WARNING: AMD & Nvidia support in this ebuild are experimental"
		einfo "If you run into issues, especially compiling dev-libs/rocm-opencl-runtime"
		einfo "you may try the docker image here https://github.com/ROCm/ROCm-docker"
		einfo "and follow instructions here"
		einfo "https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/docker.html"
	fi
}

src_unpack() {
	git-r3_src_unpack
	go-module_live_vendor
}

src_prepare() {
	default

	if use video_cards_amdgpu; then
		# --hip-version gets appended to the compile flags which isn't a known flag.
		# This causes rocm builds to fail because -Wunused-command-line-argument is turned on.
		# Use nuclear option to fix this.
	# Disable -Werror's from go modules.
		find "${S}" -name ".go" -exec sed -i "s/ -Werror / /g" {} + || die
	fi
}

src_compile() {
	CUSTOM_CPU_FLAGS=""
	use cpu_flags_x86_avx && CUSTOM_CPU_FLAGS+="avx"
	use cpu_flags_x86_avx2 && CUSTOM_CPU_FLAGS+=",avx2"
	use cpu_flags_x86_avx512 && CUSTOM_CPU_FLAGS+=",avx512,avx512vbmi,avx512vnni,avx512bf16"

	# Build basic ollama executable with cpu features built in
	export CUSTOM_CPU_FLAGS

	if use video_cards_amdgpu; then
		export HIP_ARCHS=$(get_amdgpu_flags)
		export HIP_PATH="/usr"
	else
		export OLLAMA_SKIP_ROCM_GENERATE=1
	fi

	if ! use cuda; then
		export OLLAMA_SKIP_CUDA_GENERATE=1
	fi
	emake dist
}

src_install() {
	dobin dist/linux-${ARCH}/bin/ollama

	insinto /usr/lib
	doins -r dist/linux-${ARCH}/lib/ollama

	doinitd "${FILESDIR}"/ollama
}

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

pkg_postinst() {
	einfo "Quick guide:"
	einfo "ollama serve"
	einfo "ollama run llama3:70b"
	einfo "See available models at https://ollama.com/library"
}