#!/usr/bin/env bash
# Install llama.cpp with the Entrit d2 kernel (D2NL).
#
#   bash install.sh                   first install: fetches the other files, then builds
#   ./install.sh --update             download the newest kernel, then rebuild
#   ./install.sh --backend vulkan     force one backend: cuda | vulkan | metal | hip | cpu
#
# Private repo? Set your Hugging Face token first:  export HF_TOKEN=hf_...
set -euo pipefail

HERE=$(cd "$(dirname "$0")" && pwd)
REPO_URL=${ENTRIT_REPO:-https://huggingface.co/Entrit/llama.cpp-d2nl/resolve/main}
FILES=(install.sh get-model.sh run.sh README.md LLAMA_CPP_VERSION patches/entrit-d2nl.diff)
BACKEND=auto
UPDATE=0

while [[ $# -gt 0 ]]; do
    case "$1" in
        --backend)  BACKEND=${2:?--backend needs a value}; shift 2 ;;
        --update)   UPDATE=1; shift ;;
        --fetched)  UPDATE=0; FETCHED=1; shift ;;
        -h|--help)  sed -n '2,9p' "$0"; exit 0 ;;
        *) echo "unknown option: $1 (try --help)"; exit 1 ;;
    esac
done

say()  { printf '\n\033[1m==> %s\033[0m\n' "$*"; }
fail() { printf '\n\033[31mERROR: %s\033[0m\n' "$*"; exit 1; }
OS=$(uname -s)

need() {
    command -v "$1" >/dev/null 2>&1 && return
    case "$OS" in
        Darwin) fail "'$1' is missing. Install the build tools:  xcode-select --install   and then:  brew install cmake" ;;
        *)      fail "'$1' is missing. On Ubuntu/Debian:  sudo apt install -y build-essential cmake curl patch" ;;
    esac
}
need curl

# ------------------------------------------------------------------ 1. our files
AUTH=()
[[ -n "${HF_TOKEN:-}" ]] && AUTH=(-H "Authorization: Bearer $HF_TOKEN")
missing=0
for f in "${FILES[@]}"; do [[ -f "$HERE/$f" ]] || missing=1; done
if [[ "$UPDATE" == 1 || ( "$missing" == 1 && -z "${FETCHED:-}" ) ]]; then
    say "Downloading the Entrit files"
    mkdir -p "$HERE/patches"
    for f in "${FILES[@]}"; do
        curl -fsSL ${AUTH[@]+"${AUTH[@]}"} -o "$HERE/$f.new" "$REPO_URL/$f" \
            || fail "could not download $f. If the repository is private, set HF_TOKEN (see the top of this file)."
        mv "$HERE/$f.new" "$HERE/$f"
    done
    chmod +x "$HERE/install.sh" "$HERE/get-model.sh" "$HERE/run.sh"
    exec bash "$HERE/install.sh" --fetched --backend "$BACKEND"    # run the version we just downloaded
fi

# ------------------------------------------------------------------ 2. tools
say "Checking tools"
need cmake; need c++; need tar; need patch
echo "cmake, a C++ compiler, tar and patch are present"

# ------------------------------------------------------------------ 3. backend
if [[ "$BACKEND" == auto ]]; then
    if   [[ "$OS" == Darwin ]];                                               then BACKEND=metal
    elif command -v nvcc >/dev/null 2>&1 || [[ -x /usr/local/cuda/bin/nvcc ]]; then BACKEND=cuda
    elif command -v hipcc >/dev/null 2>&1;                                     then BACKEND=hip
    elif command -v glslc >/dev/null 2>&1;                                     then BACKEND=vulkan
    else                                                                            BACKEND=cpu
    fi
fi
case "$BACKEND" in
    metal)  FLAGS=(-DGGML_METAL=ON) ;;
    cuda)   [[ -x /usr/local/cuda/bin/nvcc ]] && export PATH="/usr/local/cuda/bin:$PATH"
            FLAGS=(-DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=native) ;;   # build for the GPU in this machine only
    hip)    FLAGS=(-DGGML_HIP=ON) ;;
    vulkan) command -v glslc >/dev/null 2>&1 || fail "Vulkan needs glslc. Ubuntu: sudo apt install -y libvulkan-dev glslc   Other: install the Vulkan SDK from https://vulkan.lunarg.com"
            FLAGS=(-DGGML_VULKAN=ON) ;;
    cpu)    FLAGS=(-DGGML_CUDA=OFF -DGGML_VULKAN=OFF -DGGML_METAL=OFF) ;;
    *)      fail "unknown backend '$BACKEND' (use cuda, vulkan, metal, hip or cpu)" ;;
esac
say "Backend: $BACKEND"
[[ "$BACKEND" == cpu ]] && echo "No GPU toolkit found, so this build runs on the CPU only. It works, but it is slow."

# ------------------------------------------------------------------ 4. source
VERSION=$(tr -d '[:space:]' < "$HERE/LLAMA_CPP_VERSION")
SRC="$HERE/llama.cpp"
say "Downloading llama.cpp (version ${VERSION:0:7})"
rm -rf "$SRC" && mkdir -p "$SRC"
curl -fsSL "https://github.com/PrismML-Eng/llama.cpp/archive/$VERSION.tar.gz" | tar -xz -C "$SRC" --strip-components 1 \
    || fail "could not download llama.cpp from GitHub. Check your internet connection."

say "Adding the Entrit kernel"
POPTS=(-p1 -s)
patch --version 2>/dev/null | grep -q GNU && POPTS+=(--no-backup-if-mismatch)   # GNU patch (Linux); macOS patch has no .orig files to avoid
(cd "$SRC" && patch "${POPTS[@]}" < "$HERE/patches/entrit-d2nl.diff") \
    || fail "the kernel patch did not apply. Run ./install.sh --update and try again."
chmod +x "$SRC"/entrit/*.sh "$SRC"/entrit/*.py 2>/dev/null || true

# ------------------------------------------------------------------ 5. build
say "Building (this takes 5-20 minutes the first time)"
BUILD="$HERE/build"
JOBS=$( (command -v nproc >/dev/null && nproc) || sysctl -n hw.ncpu 2>/dev/null || echo 4)
cmake -S "$SRC" -B "$BUILD" -DCMAKE_BUILD_TYPE=Release "${FLAGS[@]}" > "$HERE/build-config.log" 2>&1 \
    || fail "configure failed; the details are in build-config.log"
cmake --build "$BUILD" -j "$JOBS" --target llama-server llama-cli llama-completion llama-bench test-backend-ops \
    > "$HERE/build.log" 2>&1 || fail "build failed; the details are in build.log"

# ------------------------------------------------------------------ 6. self-test
DEV=""
case "$BACKEND" in cuda|hip) DEV=CUDA0 ;; vulkan) DEV=Vulkan0 ;; metal) DEV=MTL0 ;; esac
if [[ -n "$DEV" ]]; then
    say "Self-test: Entrit kernel on $DEV vs the CPU"
    if "$BUILD/bin/test-backend-ops" test -o MUL_MAT -b "$DEV" -p d2nl 2>&1 | grep -q "tests passed"; then
        echo "PASS"
    else
        echo "The self-test did not pass. The build may still run, but please report this, with build.log."
    fi
fi

say "Done"
echo "Next: download a model with   ./get-model.sh <huggingface-repo>"
echo "Then run it with              ./run.sh models/<file>.gguf"