#!/usr/bin/env bash
# Build patched llama.cpp (AtomicChat ricursive port) for RTX PRO 6000 Blackwell (sm_120)
# and upload: (1) llama-server tarball for the ZeroGPU Space, (2) llama-cpp-python wheel (best effort).
set -euxo pipefail

REPO="akhaliq/llama-cpp-ricursive-cuda120"
COMMIT="9e2a04a439edc49f08c5acdafefae19dfa25a658"
PATCH_URL="https://huggingface.co/AtomicChat/Ripple-65B-Preview-GGUF/resolve/main/ricursive-port.patch"
WORK=/tmp/work
mkdir -p "$WORK" && cd "$WORK"

export DEBIAN_FRONTEND=noninteractive
apt-get update -qq
apt-get install -y -qq git cmake ninja-build curl python3 python3-pip
pip3 install -q -U huggingface_hub scikit-build-core

# ---------------- patched llama.cpp source ----------------
git clone -q https://github.com/ggml-org/llama.cpp
cd llama.cpp
git checkout -q "$COMMIT"
curl -fsSL -o /tmp/ricursive-port.patch "$PATCH_URL"
git apply /tmp/ricursive-port.patch
echo "patch applied cleanly"

# ---------------- build llama-server (CUDA sm_120) ----------------
CMAKE_FLAGS=(
  -G Ninja
  -DCMAKE_BUILD_TYPE=Release
  -DLLAMA_CURL=OFF
  -DGGML_CUDA=ON
  -DCMAKE_CUDA_ARCHITECTURES=120
  -DGGML_NATIVE=OFF
  -DGGML_BACKEND_DL=ON
  -DGGML_CPU_ALL_VARIANTS=ON
)
if ! cmake -B build "${CMAKE_FLAGS[@]}"; then
  echo "cmake with CPU_ALL_VARIANTS failed; retrying without it"
  rm -rf build
  cmake -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLAMA_CURL=OFF \
    -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=120 -DGGML_NATIVE=OFF
fi
cmake --build build -j"$(nproc)" --target llama-server

# Stage CUDA runtime libs next to the binaries: the ZeroGPU image has no CUDA toolkit,
# and libggml-cuda.so silently fails to dlopen without them (CPU fallback).
echo "=== ldd before staging ==="
ldd build/bin/llama-server || true
ldd build/bin/libggml-cuda.so || true
for lib in $(ldd build/bin/*.so build/bin/llama-server 2>/dev/null | grep -oP '/usr/local/cuda/lib64/\S+' | sort -u); do
  base=$(basename "$lib")
  case "$base" in
    libcuda.so*) echo "skip driver lib $base (provided by the GPU host)";;
    *) cp -L "$lib" build/bin/ && echo "staged $base";;
  esac
done
echo "=== ldd after staging (LD_LIBRARY_PATH=build/bin) ==="
LD_LIBRARY_PATH=build/bin ldd build/bin/libggml-cuda.so | grep -E "not found|cuda" || true
ls -la build/bin

tar -C build -czf /tmp/llama-ricursive-cuda128-sm120.tar.gz bin
ls -la /tmp/llama-ricursive-cuda128-sm120.tar.gz

python3 - <<'EOF'
from huggingface_hub import HfApi
HfApi().upload_file(
    path_or_fileobj="/tmp/llama-ricursive-cuda128-sm120.tar.gz",
    path_in_repo="llama-ricursive-cuda128-sm120.tar.gz",
    repo_id="akhaliq/llama-cpp-ricursive-cuda120",
    repo_type="model",
    commit_message="llama-server: ricursive port (AtomicChat) @ 9e2a04a, CUDA 12.8 sm_120",
)
print("tarball uploaded")
EOF

# ---------------- llama-cpp-python wheel (best effort, non-fatal) ----------------
set +e
cd "$WORK"
git clone -q --depth 1 https://github.com/abetlen/llama-cpp-python
cd llama-cpp-python
rm -rf vendor/llama.cpp
mkdir -p vendor
cp -r "$WORK/llama.cpp" vendor/llama.cpp
rm -rf vendor/llama.cpp/.git vendor/llama.cpp/build

export CMAKE_ARGS="-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=120 -DGGML_NATIVE=OFF -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DLLAMA_CURL=OFF"
pip3 wheel . --no-deps -w /tmp/wheelhouse
WHEEL_STATUS=$?
set -e
if [ $WHEEL_STATUS -eq 0 ]; then
  cd /tmp/wheelhouse
  for w in llama_cpp_python-*.whl; do
    new=$(echo "$w" | sed 's/llama_cpp_python-\([0-9.]*\)-/llama_cpp_python-\1+cu128.ricursive-/')
    mv "$w" "$new"
  done
  ls -la /tmp/wheelhouse
  python3 - <<'EOF'
import glob
from huggingface_hub import HfApi
for whl in glob.glob("/tmp/wheelhouse/*.whl"):
    name = whl.rsplit("/", 1)[-1]
    HfApi().upload_file(
        path_or_fileobj=whl,
        path_in_repo=name,
        repo_id="akhaliq/llama-cpp-ricursive-cuda120",
        repo_type="model",
        commit_message="llama-cpp-python wheel: ricursive port, CUDA 12.8 sm_120",
    )
    print("wheel uploaded:", name)
EOF
else
  echo "wheel build failed (non-fatal); tarball is already uploaded"
fi

echo "BUILD COMPLETE"