#!/usr/bin/env bash
# download-qwen38-bf16.sh — RUNS ON DANKSERVER. Fetch and verify the vendor
# BF16 checkpoint Qwen/Qwen3.8-Flash-Next: the source the five-bit campaign
# converts, and the source the KLD reference is recorded from.
#
# Why the base and not a finetune: the campaign converts to compare against the
# publisher's own iq4r container, and that container was built from this
# checkpoint. Base-in, base-out; anything else confounds the gate. (The
# predecessor campaign ran ukisai/Swift-1.5-Qwen3.8-Flash-Next-BF16, whose key
# set is identical to this repo's but whose weights are not — none of its
# calibration grams, KLD reference or protected list carries over.)
#
# Why re-fetch instead of reusing: every BF16 copy of this family was destroyed
# in the 2026-10-15 storage purge — only the 113.6 GiB served .rad survived.
# 335.3 GiB of safetensors, 131 shards, ungated (the predecessor campaign
# verified the key set against this repo, so no gate letter is needed).
#
# Verification is four gates deep, because the failure this repo actually lived
# through is a fetch that reported success: `curl -o FILE URL` exits 0 on an
# HTTP 401/403 and writes the error body as the artifact — three shards of
# identical tiny size and a cheerful DONE in 20 seconds for 104 GiB.
#
#   1. shard census: the file count must equal the expected shard count, so a
#      partially fetched fragment tree (the .partNofM trap) cannot pass.
#   2. per-shard completeness, from the shard's OWN safetensors header: the
#      header declares where the tensor data ends, so every shard's true size is
#      known offline. A torn tail is named here, by file, instead of surfacing
#      as a load error five hours later.
#   3. index coverage: every shard model.safetensors.index.json references
#      exists and is named.
#   4. total bytes against the number the Hub declared, then, unless --no-hub,
#      every file's sha256 against the Hub's LFS record — including the shards
#      already on disk, so a previous partial download is never trusted.
#
# The token never enters a command line (the 2026-09-16 leak law): it is read
# into HF_TOKEN in the environment, where a local process listing is the only
# reader. This repo is ungated, so the fetch simply runs without one when no
# token file exists. scripts/hfget.sh is curl+jq only — no Python, no venv, no
# hf_transfer — and it drops the pickle/ONNX weight duplicates by itself.
#
# usage: download-qwen38-bf16.sh [--fetch] [--no-hub] [--dir DIR]
#                                [--shards N] [--bytes B] [--repo ID]
#   --fetch    download first (resumable: re-run to continue), then verify
#   --no-hub   skip the Hub-side sha256 sweep; the other gates always run
# exit: 0 verified, 1 verification failed, 2 environment / usage error
set -euo pipefail

REPO="${FRANK_QWEN38_HF_REPO:-Qwen/Qwen3.8-Flash-Next}"
LOCAL_DIR="${FRANK_QWEN38_CKPT:-/home/dank/models/Qwen3.8-Flash-Next}"
# Hub census of 2026-10-27: 144 files, 131 of them .safetensors, together
# 360,023,351,514 bytes; the shards alone are 360,000,192,888. This gate counts
# shards only, so a config file edited on the Hub cannot fail it.
EXPECTED_SHARDS="${FRANK_QWEN38_SHARDS:-131}"
EXPECTED_BYTES="${FRANK_QWEN38_BYTES:-360000192888}"
MIN_FREE_MIB="${FRANK_QWEN38_MIN_FREE_MIB:-350000}"

FETCH=0
HUB=1
while [ $# -gt 0 ]; do
    case "$1" in
        --fetch) FETCH=1 ;;
        --no-hub) HUB=0 ;;
        --dir)    LOCAL_DIR="$2"; shift ;;
        --shards) EXPECTED_SHARDS="$2"; shift ;;
        --bytes)  EXPECTED_BYTES="$2"; shift ;;
        --repo)   REPO="$2"; shift ;;
        *) echo "usage: ${0##*/} [--fetch] [--no-hub] [--dir DIR] [--shards N] [--bytes B] [--repo ID]" >&2
           exit 2 ;;
    esac
    shift
done

SELF_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FAILS=0
note() { echo "$@"; FAILS=$((FAILS + 1)); }

if [ "$FETCH" = "1" ]; then
    # Every fetch gate is a hard stop: skipping the download silently and then
    # verifying a half tree is how a campaign ships a stale payload.
    HFGET="$SELF_DIR/hfget.sh"
    if [ ! -f "$HFGET" ]; then
        echo "REFUSED: the fetch helper is not here: $HFGET" >&2
        exit 2
    fi
    if ! command -v jq >/dev/null 2>&1; then
        echo "REFUSED: hfget.sh needs jq on PATH" >&2
        exit 2
    fi
    mkdir -p "$LOCAL_DIR"
    # The campaign's peak-disk law starts at the fetch: the checkpoint and the
    # reference container built from it are ~335 GiB each, and a tree that does
    # not fit is discovered hours into a transfer, when it costs the most.
    free_mib=$(df -m --output=avail "$LOCAL_DIR" | tail -1 | tr -dc '0-9')
    if [ "${free_mib:-0}" -lt "$MIN_FREE_MIB" ]; then
        echo "REFUSED: $free_mib MiB free at $LOCAL_DIR; the fetch wants $MIN_FREE_MIB MiB free" >&2
        echo "         (retire the peak-disk predecessors first — nothing is downloaded, so" >&2
        echo "         nothing is left half-written)" >&2
        exit 2
    fi
    echo "fetch: $REPO -> $LOCAL_DIR ($free_mib MiB free)"
    if [ -z "${HF_TOKEN:-}" ]; then
        CREDSTORE="$HOME/.config/huggingface/credentials/https://huggingface.co"
        if [ -f "$CREDSTORE" ]; then
            HF_TOKEN="$(awk '/^ OauthToken: / {print $2; exit}' "$CREDSTORE")"
        elif [ -f "$HOME/.cache/huggingface/token" ]; then
            # The legacy single-line token file the hub library writes.
            HF_TOKEN="$(tr -d '[:space:]' < "$HOME/.cache/huggingface/token")"
        elif [ -f "$HOME/.hf_token" ]; then
            HF_TOKEN="$(tr -d '[:space:]' < "$HOME/.hf_token")"
        fi
    fi
    if [ -n "${HF_TOKEN:-}" ]; then
        export HF_TOKEN
    fi
    bash "$HFGET" "$REPO" "$LOCAL_DIR" main
    unset HF_TOKEN
fi

if [ ! -d "$LOCAL_DIR" ]; then
    echo "REFUSED: $LOCAL_DIR does not exist (run with --fetch, or pass --dir)" >&2
    exit 1
fi
INDEX="$LOCAL_DIR/model.safetensors.index.json"
if [ ! -f "$INDEX" ]; then
    echo "REFUSED: $INDEX is missing — the checkpoint tree is not addressable without it" >&2
    exit 1
fi

echo "== gate 1: shard census ($LOCAL_DIR)"
SHARDS=$(find "$LOCAL_DIR" -maxdepth 1 -type f -name '*.safetensors' | wc -l)
if [ "$SHARDS" != "$EXPECTED_SHARDS" ]; then
    note "FAIL SHARD_COUNT: $SHARDS safetensors files on disk, expected $EXPECTED_SHARDS"
else
    echo "  ok: $SHARDS shards"
fi

echo "== gate 2: every shard is whole, judged by its own header"
# A safetensors file is [u64 header length][JSON header][tensor data]. The JSON
# declares where the data ends, so the file's true size is computable offline:
# 8 + header + the largest end offset. Truncation and error bodies both fall
# here, and each is reported with the file that caused it.
while IFS= read -r f; do
    state=$(python3 - "$f" <<'PY'
import json, os, struct, sys
path = sys.argv[1]
size = os.path.getsize(path)
try:
    with open(path, "rb") as fh:
        (hlen,) = struct.unpack("<Q", fh.read(8))
        if hlen == 0 or hlen > size:
            print("NOT_SAFETENSORS")
            raise SystemExit
        header = json.loads(fh.read(hlen))
except Exception:
    print("NOT_SAFETENSORS")
    raise SystemExit
if not header:
    print("NOT_SAFETENSORS")
    raise SystemExit
end = max((v["data_offsets"][1] for v in header.values()
           if isinstance(v, dict) and "data_offsets" in v), default=0)
want = 8 + hlen + end
print("OK" if want == size else f"TRUNCATED want={want} have={size}")
PY
)
    case "$state" in
        OK) ;;
        NOT_SAFETENSORS*) note "FAIL NOT_SAFETENSORS: $(basename "$f") is not a safetensors file (an error body written as an artifact?)" ;;
        TRUNCATED*) note "FAIL TRUNCATED: $(basename "$f") $state" ;;
    esac
done < <(find "$LOCAL_DIR" -maxdepth 1 -type f -name '*.safetensors' | sort)
if [ "$FAILS" -eq 0 ]; then
    echo "  ok: every shard's data extent matches its size on disk"
fi

echo "== gate 3: the index's shards are all present"
INDEX_SHARDS=$(python3 - "$INDEX" <<'PY'
import json, sys
map_ = json.load(open(sys.argv[1]))["weight_map"]
print("\n".join(sorted(set(map_.values()))))
PY
)
while IFS= read -r name; do
    [ -n "$name" ] || continue
    if [ ! -f "$LOCAL_DIR/$name" ]; then
        note "FAIL MISSING: $name is named by the index and is not on disk"
    fi
done <<< "$INDEX_SHARDS"
if [ "$FAILS" -eq 0 ]; then
    echo "  ok: $(printf '%s\n' "$INDEX_SHARDS" | grep -c .) referenced shards on disk"
fi

echo "== gate 4: total bytes vs the Hub's declaration"
TOTAL=$(find "$LOCAL_DIR" -maxdepth 1 -type f -name '*.safetensors' -printf '%s\n' | awk '{s+=$1} END{print s+0}')
if [ "$TOTAL" != "$EXPECTED_BYTES" ]; then
    note "FAIL TOTAL_BYTES: shards total $TOTAL bytes, expected $EXPECTED_BYTES"
else
    echo "  ok: $TOTAL bytes"
fi

if [ "$HUB" = "1" ]; then
    echo "== gate 5: every file's sha256 vs the Hub's LFS record"
    # size = file size; lfs.sha256 = sha256 of the file itself (LFS or not).
    TMPD="$(mktemp -d)"
    trap 'rm -rf "$TMPD"' EXIT
    MANIFEST="$TMPD/tree.json"
    PLAN="$TMPD/plan"
    curl -sf "https://huggingface.co/api/models/$REPO/tree/main?recursive=true&expand=true" \
        -o "$MANIFEST" || { echo "FAIL: the Hub API call failed — refusing to report success" >&2; exit 1; }
    python3 - "$MANIFEST" > "$PLAN" <<'PY'
import json, sys
for e in json.load(open(sys.argv[1])):
    if e.get("type") != "file":
        continue
    sha = (e.get("lfs") or {}).get("sha256") or e.get("sha") or ""
    if sha:
        print(e["path"], e["size"], sha)
PY
    BAD=0
    while read -r rel size sha; do
        [ -n "${rel:-}" ] || continue
        if [ ! -f "$LOCAL_DIR/$rel" ]; then
            note "FAIL MISSING: $rel is on the Hub and not on disk"
            continue
        fi
        have=$(stat -c %s "$LOCAL_DIR/$rel")
        if [ "$have" != "$size" ]; then
            note "FAIL SIZE: $rel is $have bytes, Hub says $size"
            continue
        fi
        got=$(sha256sum "$LOCAL_DIR/$rel" | awk '{print $1}')
        if [ "$got" != "$sha" ]; then
            note "FAIL SHA256: $rel got $got want $sha"
            BAD=$((BAD + 1))
        fi
    done < "$PLAN"
    if [ "$BAD" -eq 0 ] && [ "$FAILS" -eq 0 ]; then
        echo "  ok: sha256 matches for every file"
    fi
    rm -rf "$TMPD"
    trap - EXIT
fi

if [ "$FAILS" -gt 0 ]; then
    echo "VERIFY_FAILED: $FAILS problem(s) in $LOCAL_DIR" >&2
    exit 1
fi
echo "VERIFY_OK: $SHARDS shards, $TOTAL bytes of safetensors in $LOCAL_DIR"
echo "NEXT: build the BF16 reference container with scripts/build-qwen38-bf16-rad.sh"