#!/usr/bin/env bash
# Thanatos-27B — heal a previously pulled HF-bridge tag whose bundled
# GGUF is `qwen36`-stamped (legacy v0.6.0-era pulls before `964e418`,
# 3rd-round-trip-era pulls between `973d7ef` and `978798f`, or
# 5th-round-trip-era pulls between `ae67ed1` and `e03e10e`).
#
# Fresh pulls of `ollama run hf.co/FoolDev/Thanatos-27B` now get the
# qwen35-stamped bundle and load directly — this script is the
# recovery path for users who pulled a qwen36-stamped blob into
# their local Ollama store during one of the qwen36 windows
# and haven't refreshed since.
#
# It rebadges the HF-bridge tag's model blob in-place (qwen36 ->
# qwen35, metadata-only, byte-identical tensors) and rewrites the
# manifest's model-layer digest to point at the new blob. After
# running, the cached `hf.co/FoolDev/Thanatos-27B` tag loads.
#
# Idempotent: a tag already on qwen35 / qwen35moe is left untouched.
# The current bundle is qwen35-stamped so this script is a no-op for
# anyone who pulled after the latest re-stamp; it stays in the repo
# for the legacy recovery case.
#
# Usage:
#   ./scripts/heal_hf_pull.sh                                # default tag
#   TAG=hf.co/FoolDev/Thanatos-27B:Q4_K_M ./scripts/heal_hf_pull.sh
#
# Requires: ollama, jq, python3 with the `gguf` package, sha256sum.
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
TAG="${TAG:-hf.co/FoolDev/Thanatos-27B:Q4_K_M}"
OLLAMA_MODELS="${OLLAMA_MODELS:-${HOME}/.ollama/models}"

red()   { printf "\033[31m%s\033[0m\n" "$*"; }
green() { printf "\033[32m%s\033[0m\n" "$*"; }
blue()  { printf "\033[34m%s\033[0m\n" "$*"; }

blue "[*] tag:    ${TAG}"
blue "[*] store:  ${OLLAMA_MODELS}"

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

for bin in ollama jq python3 sha256sum; do
    if ! command -v "${bin}" >/dev/null 2>&1; then
        red "[!] missing dependency: ${bin}"; exit 1
    fi
done

# ---- 2. Locate the model blob and manifest ----------------------------------

# `ollama show --modelfile` writes a FROM line with the absolute blob path.
# Reliable regardless of which case variant the user pulled with
# (hf.co's 307 lets `Thanatos-27B` and `thanatos-27b` both resolve to the
# canonical repo, and ollama stores the manifest under whichever case
# was first registered).
#
# The `if MODELFILE_OUT=...; then` form rather than a direct command
# substitution is load-bearing: under `set -e + pipefail`, a bare
# `MODEL_BLOB="$(ollama show ... | awk ...)"` silently terminates
# the script when the tag isn't pulled (ollama show exits non-zero
# -> pipefail propagates -> set -e exits the script before the
# explicit `[[ -z "${MODEL_BLOB}" ]]` check below ever runs). The
# `if` form takes the false branch on failure without tripping
# set -e, so the user gets the actionable error instead of a silent
# `make: *** [Makefile:N: heal-hf] Error 1`.
MODEL_BLOB=""
if MODELFILE_OUT="$(ollama show --modelfile "${TAG}" 2>/dev/null)"; then
    MODEL_BLOB="$(awk '/^FROM[[:space:]]/ {print $2; exit}' <<<"${MODELFILE_OUT}")"
fi
if [[ -z "${MODEL_BLOB}" || ! -f "${MODEL_BLOB}" ]]; then
    red "[!] could not resolve model blob for tag '${TAG}'."
    red "    Is the tag pulled? Try: ollama pull ${TAG}"
    exit 1
fi
MODEL_HASH="$(basename "${MODEL_BLOB}" | sed 's/^sha256-//')"
blue "[*] blob:   ${MODEL_BLOB}"

# Find the manifest by grepping for the model digest. The blob is
# referenced from exactly one tag in the heal scenario — fresh HF pull
# of a single :Q4_K_M tag — but if someone has multiple tags pointing
# at the same blob, we filter down to the one matching ${TAG}.
TAG_PATH="${TAG#hf.co/}"      # FoolDev/Thanatos-27B:Q4_K_M
NAMESPACE_PATH="${TAG_PATH%:*}" # FoolDev/Thanatos-27B
TAG_FILE="${TAG_PATH##*:}"    # Q4_K_M

MANIFEST="$(find "${OLLAMA_MODELS}/manifests/hf.co" \
              -type f \
              -ipath "*/${NAMESPACE_PATH}/${TAG_FILE}" 2>/dev/null | head -1)"

if [[ -z "${MANIFEST}" || ! -f "${MANIFEST}" ]]; then
    red "[!] manifest not found under ${OLLAMA_MODELS}/manifests/hf.co for tag '${TAG}'."
    exit 1
fi
blue "[*] manifest: ${MANIFEST}"

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

ARCH="$(python3 - "${MODEL_BLOB}" <<'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
)"
blue "[*] arch:    ${ARCH}"

if [[ "${ARCH}" == "qwen35" || "${ARCH}" == "qwen35moe" ]]; then
    green "[=] already on a loadable arch (${ARCH}) — nothing to heal."
    exit 0
fi
if [[ "${ARCH}" != "qwen36" ]]; then
    red "[!] unexpected arch '${ARCH}' — refusing to heal. Edit this script if intentional."
    exit 1
fi

# ---- 4. Rebadge to a temp blob and stage it in the store --------------------

# Stage in the repo's .cache/ rather than /tmp: the rebadged copy is the same
# size as the original (~17 GB), which blows past a typical tmpfs /tmp budget.
# .cache/ is on the same filesystem as ~/.ollama on a normal Linux home dir
# layout, so the final move into blobs/ is an atomic rename, not a copy.
SCRATCH_DIR="${ROOT}/.cache"
mkdir -p "${SCRATCH_DIR}"
TMP_BLOB="$(mktemp -p "${SCRATCH_DIR}" thanatos-heal.XXXXXX.gguf)"
trap 'rm -f "${TMP_BLOB}"' EXIT
blue "[*] rebadging qwen36 -> qwen35 (metadata only, tensors byte-identical) ..."
python3 "${ROOT}/scripts/rename_arch.py" \
    --from-arch qwen36 --to-arch qwen35 \
    "${MODEL_BLOB}" "${TMP_BLOB}"

NEW_HASH="$(sha256sum "${TMP_BLOB}" | awk '{print $1}')"
NEW_SIZE="$(stat -c '%s' "${TMP_BLOB}")"
NEW_BLOB="${OLLAMA_MODELS}/blobs/sha256-${NEW_HASH}"
blue "[*] new digest: sha256:${NEW_HASH}"
blue "[*] new size:   ${NEW_SIZE}"

if [[ -f "${NEW_BLOB}" ]]; then
    blue "[=] target blob already in store — reusing."
    rm -f "${TMP_BLOB}"
else
    mv "${TMP_BLOB}" "${NEW_BLOB}"
fi
trap - EXIT

# ---- 5. Rewrite the manifest's model layer ----------------------------------

# Keep a sidecar copy of the original manifest so we can roll back if
# `ollama show` rejects the rewrite. The blob cleanup in step 7 only
# happens AFTER validation passes, so a rollback always lands on a
# consistent (original manifest, original blob still present) state.
BACKUP_MANIFEST="${MANIFEST}.heal-backup"
cp "${MANIFEST}" "${BACKUP_MANIFEST}"

TMP_MANIFEST="$(mktemp -t thanatos-heal-manifest.XXXXXX.json)"
trap 'rm -f "${TMP_MANIFEST}"' EXIT
jq --arg new "sha256:${NEW_HASH}" \
   --argjson size "${NEW_SIZE}" '
    .layers |= map(
        if .mediaType == "application/vnd.ollama.image.model"
        then .digest = $new | .size = $size
        else .
        end
    )
' "${MANIFEST}" > "${TMP_MANIFEST}"

NEW_DIGEST_IN_MANIFEST="$(jq -r '
    .layers[] | select(.mediaType == "application/vnd.ollama.image.model") | .digest
' "${TMP_MANIFEST}")"
if [[ "${NEW_DIGEST_IN_MANIFEST}" != "sha256:${NEW_HASH}" ]]; then
    red "[!] manifest rewrite failed (digest mismatch); not committing."
    rm -f "${BACKUP_MANIFEST}"
    exit 1
fi
mv "${TMP_MANIFEST}" "${MANIFEST}"
trap - EXIT

# ---- 6. Validate the rewritten manifest via ollama show ---------------------

# `ollama show` parses the manifest, walks the layer digests, and reads
# enough of the model blob to extract metadata (arch, params, etc.).
# A buggy rewrite that produces jq-accepted JSON but breaks an ollama
# invariant (layer order, mediaType set, digest-vs-blob-bytes
# mismatch) trips here, before we lose the rollback path by removing
# the old qwen36 blob. On failure we restore from BACKUP_MANIFEST.
blue "[*] validating rewritten manifest with 'ollama show'..."
if ! ollama show "${TAG}" >/dev/null 2>&1; then
    red "[!] ollama show ${TAG} failed after the manifest rewrite."
    blue "[*] rolling back: restoring original manifest from ${BACKUP_MANIFEST}"
    mv "${BACKUP_MANIFEST}" "${MANIFEST}"
    red "    Tag is back to its pre-heal (qwen36) state. The new qwen35"
    red "    blob is left in the store; ollama auto-prunes unreferenced"
    red "    blobs on next restart."
    exit 1
fi
rm -f "${BACKUP_MANIFEST}"
green "[+] manifest validates."

# ---- 7. Remove the old qwen36 blob if no other manifest references it -------

OLD_DIGEST="sha256:${MODEL_HASH}"
if ! grep -rlF -- "${OLD_DIGEST}" "${OLLAMA_MODELS}/manifests/" >/dev/null 2>&1; then
    blue "[*] no other manifest references the old qwen36 blob — removing ${MODEL_BLOB}"
    rm -f "${MODEL_BLOB}"
else
    blue "[=] old qwen36 blob still referenced by another manifest — leaving in place."
fi

echo
green "[+] healed. Try it:"
echo "    ollama run ${TAG}"
echo "    MODEL=${TAG} make smoke"