#!/usr/bin/env bash
# Janus-35B — load this repo's bundle into Ollama as a local tag.
#
# The bundled GGUF (Janus-35B-A3B.Q4_K_M.gguf) is qwen35moe-stamped and
# loads directly on stock llama.cpp / Ollama. This script is the
# one-shot path from "I just cloned this repo" to "I have a working
# local Ollama tag":
#
#   1. Resolve the bundle. If it's an LFS pointer (cloned without
#      `git lfs pull`), download the real ~19 GB blob via `hf download`.
#   2. Sanity-check `general.architecture` is qwen35moe / qwen35.
#   3. Run `ollama create <tag> -f <temp Modelfile pointing at the
#      resolved bundle>`.
#
# Useful if you want a bare local tag (`janus`) rather than the
# `hf.co/FoolDev/Janus-35B-HERETIC` path.
#
# Usage:
#   ./scripts/load_bundle.sh                 # default tag: janus
#   TAG=janus-bundle ./scripts/load_bundle.sh
#   BUNDLE=/path/to/Janus-35B-A3B.Q4_K_M.gguf ./scripts/load_bundle.sh
#
# Requires: ollama, python3 with the `gguf` package, hf (if the bundle
# needs to be downloaded).
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BUNDLE="${BUNDLE:-${ROOT}/Janus-35B-A3B.Q4_K_M.gguf}"
TAG="${TAG:-janus}"
REPO_ID="${REPO_ID:-FoolDev/Janus-35B-HERETIC}"
MODELFILE="${ROOT}/Modelfile"

echo "[*] bundle:    ${BUNDLE}"
echo "[*] tag:       ${TAG}"

# ---- 1. Sanity ---------------------------------------------------------------

if ! command -v ollama >/dev/null 2>&1; then
    echo "[!] ollama not found in PATH" >&2; exit 1
fi
if [[ ! -f "${MODELFILE}" ]]; then
    echo "[!] missing ${MODELFILE}" >&2; exit 1
fi

# ---- 2. Resolve bundle (smudge LFS pointer if needed) ------------------------

resolve_bundle() {
    local file="$1"
    if [[ ! -f "${file}" ]]; then
        return 1
    fi
    # LFS pointer files are tiny (a couple hundred bytes) and start with
    # `version https://git-lfs.github.com/spec/v1`.
    local size
    size="$(stat -c '%s' "${file}")"
    if (( size < 1024 )) && head -n1 "${file}" | grep -q 'git-lfs'; then
        return 1
    fi
    return 0
}

if ! resolve_bundle "${BUNDLE}"; then
    # Download to a side path under .cache/ so we don't overwrite the
    # LFS pointer in the working tree. Without git-lfs installed, the
    # pointer never auto-smudges and the user expects the file in the
    # repo root to stay a few hundred bytes. Downstream steps read
    # whichever path BUNDLE points at, so just re-point it here.
    CACHE_DIR="${ROOT}/.cache"
    BUNDLE_NAME="$(basename "${BUNDLE}")"
    CACHED="${CACHE_DIR}/${BUNDLE_NAME}"
    if resolve_bundle "${CACHED}"; then
        echo "[=] using previously downloaded bundle at ${CACHED}"
        BUNDLE="${CACHED}"
    else
        echo "[*] bundle missing or LFS-pointer-only — downloading from ${REPO_ID} to ${CACHED} ..."
        HF=""
        if command -v hf >/dev/null 2>&1; then
            HF="hf"
        elif command -v huggingface-cli >/dev/null 2>&1; then
            HF="huggingface-cli"
        else
            echo "[!] neither 'hf' nor 'huggingface-cli' installed; can't fetch bundle" >&2
            echo "    pip install -U huggingface_hub" >&2
            exit 1
        fi
        mkdir -p "${CACHE_DIR}"
        case "${HF}" in
            hf)              hf download "${REPO_ID}" "${BUNDLE_NAME}" --local-dir "${CACHE_DIR}" ;;
            huggingface-cli) huggingface-cli download "${REPO_ID}" "${BUNDLE_NAME}" --local-dir "${CACHE_DIR}" ;;
        esac
        BUNDLE="${CACHED}"
    fi
    if ! resolve_bundle "${BUNDLE}"; then
        echo "[!] still no usable bundle at ${BUNDLE} after download" >&2; exit 1
    fi
fi

# ---- 3. Inspect arch ---------------------------------------------------------

ARCH="$(python3 - "${BUNDLE}" <<'PY'
import sys
from gguf import GGUFReader, constants
r = GGUFReader(sys.argv[1], "r")
f = r.get_field(constants.Keys.General.ARCHITECTURE)
print(bytes(f.parts[f.data[0]]).decode())
PY
)"
echo "[*] bundle arch: ${ARCH}"

if [[ "${ARCH}" != "qwen35" && "${ARCH}" != "qwen35moe" ]]; then
    echo "[!] unexpected arch '${ARCH}' — refusing to load. Edit this script if intentional." >&2
    exit 1
fi

# ---- 4. Build a Modelfile copy with FROM pointing at the bundle --------------

TMP_MODELFILE="$(mktemp -t janus35b-loadbundle.XXXXXX)"
trap 'rm -f "${TMP_MODELFILE}"' EXIT
awk -v p="${BUNDLE}" '
    /^FROM[[:space:]]/ && !done { print "FROM " p; done=1; next }
    { print }
' "${MODELFILE}" > "${TMP_MODELFILE}"

# ---- 5. Create the Ollama model ----------------------------------------------

echo "[*] ollama create ${TAG} -f <patched modelfile pointing at ${BUNDLE}>"
ollama create "${TAG}" -f "${TMP_MODELFILE}"

echo
echo "[+] Done. Try it:"
echo "    ollama run ${TAG}"