#!/usr/bin/env bash
# build-frank-archon-rad.sh — build frank-archon's .rad container from the
# OFFICIAL Qwen/Qwen3.8-Flash-Next-FP8 checkpoint, at NATIVE precision.
#
#   bash scripts/build-frank-archon-rad.sh                # convert
#   bash scripts/build-frank-archon-rad.sh --plan-only    # declare, no I/O
#   bash scripts/build-frank-archon-rad.sh --restage      # redo staging
#   bash scripts/build-frank-archon-rad.sh --force        # rebuild existing
#
# WHY A STAGING DIR (the one non-obvious step): rad-convert selects its
# architecture plugin by the checkpoint's quantisation DESCRIPTOR —
# derived from config.json's quantization_config (core/format/checkpoint.cpp:
# "fp8" here). No shipped plugin claims (qwen4exp, "fp8"): the shipped
# qwen4exp plugin claims (qwen4exp, ""), the descriptor recipe-built
# containers carry (they are built from the bf16 upstream, which has no
# quantization_config at all). The match is exact with no fallback, so an
# unmodified convert refuses before declare. The descriptor is only a
# selection key; per-weight encodings are data — the checkpoint reader
# derives them from the tensors themselves (X.weight_scale_inv is the
# scale plane; checkpoint.h). Staging drops the quantization_config key
# (everything else is a relative symlink into the pristine download) so
# the convert selects the shipped plugin. EVERY weight is then copied through
# in its own encoding — block-FP8 codes + bf16 scale planes stay block-FP8 —
# except TWO value-preserving reshapes that stack-qwen38-official-fp8.py
# performs on the staging dir first, and ONE recipe rule that finishes the
# PLE table:
#   * experts: the plugin maps the STACKED layout (mlp.experts.gate_up_proj
#     [E,2F,H] with the _scale_inv plane at the same index) while the official
#     repo ships per-expert unfused tensors — staging byte-stacks them
#     (gate rows first, no requantisation).
#   * PLE: the plugin's fp8 branch wants the table codes beside a ONE-ELEMENT
#     bf16 scale plane; the official repo ships 128 SCALELESS fp8 shards plus
#     one bf16 scalar. Staging decodes to bf16 with the checkpoint's own
#     scale, and the convert's rtn rule re-encodes at that same scale —
#     every original code is restored exactly and the scale plane lands in
#     the form the plugin declares.
#
# WHY A CONTAINER AT ALL: a checkpoint dir has no mapped/file tiers, and a
# weight only a host kernel reads (the n-gram / PLE table, 51B params)
# cannot be served from one (GUIDE §3.4). rad-convert also bakes the
# tokenizer and the checkpoint's chat template; the runtime replaces the
# template with froggeric's (--override-chat-template).
#
# The convert runs rad-convert inside the digest-pinned radiance image
# (the host has no ROCm and needs none): /home/dank/models is mounted at
# /models so the staging dir's RELATIVE symlinks resolve inside the
# container. Output lands beside the raw download, owned by dank
# (rootless podman maps container-root to the invoking user).
set -euo pipefail

SRC="/home/dank/models/Qwen/Qwen3.8-Flash-Next-FP8"
STAGE="/home/dank/models/stage-qwen4exp-fp8"
OUT_DIR="/home/dank/models/rad"
OUT="$OUT_DIR/qwen4exp-fp8-native.rad"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck disable=SC1091
source "$SCRIPT_DIR/frank-archon.env"
IMAGE="$FRANK_ARCHON_IMAGE"

plan_only=0 restage=0 force=0
for a in "$@"; do
    case "$a" in
        --plan-only) plan_only=1 ;;
        --restage)   restage=1 ;;
        --force)     force=1 ;;
        *) echo "unknown flag: $a" >&2; exit 2 ;;
    esac
done

if [ ! -f "$SRC/config.json" ]; then
    echo "ERROR: raw checkpoint missing: $SRC (run scripts/download-frank-archon.sh --fetch)" >&2
    exit 1
fi
echo "verifying raw checkpoint first..."
bash "$SCRIPT_DIR/download-frank-archon.sh"

if [ -f "$OUT" ] && [ "$plan_only" -eq 0 ] && [ "$force" -eq 0 ]; then
    echo "container already exists: $OUT"
    echo "rebuild it with --force (the old file moves to ~/.AGENTTRASH)"
    exit 0
fi

# ---- staging ----------------------------------------------------------------
if [ "$restage" -eq 1 ] && [ -e "$STAGE" ]; then
    TRASH="$HOME/.AGENTTRASH/$(date +%Y%m%d)-frank-archon-stage"
    mkdir -p "$TRASH"
    mv "$STAGE" "$TRASH/stage-$(date +%H%M%S)"
    echo "old staging dir moved to $TRASH"
fi
if [ ! -f "$STAGE/config.json" ]; then
    mkdir -p "$STAGE"
    python3 - "$SRC/config.json" "$STAGE/config.json" <<'PY'
import json, sys
src, dst = sys.argv[1], sys.argv[2]
cfg = json.load(open(src))
dropped = cfg.pop("quantization_config", None)
json.dump(cfg, open(dst, "w"), indent=2)
print(f"staging config.json written (quantization_config dropped: {dropped is not None})")
PY
    # The index is COPIED like config.json, never symlinked: the reshape step
    # rewrites the stage's index, and a symlink would write through to the
    # raw checkpoint (measured 2026-10-05: the raw index was restamped with
    # the stacked layout and the raw verify failed 143/144). Shards stay
    # relative symlinks — relative so the targets resolve inside the convert
    # container (the stage's parent is mounted at /models).
    for f in "$SRC"/*; do
        name="$(basename "$f")"
        case "$name" in config.json|model.safetensors.index.json) continue ;; esac
        ln -sfn "../Qwen/Qwen3.8-Flash-Next-FP8/$name" "$STAGE/$name"
    done
    cp "$SRC/model.safetensors.index.json" "$STAGE/"
    echo "staging dir ready: $STAGE ($(ls "$STAGE" | wc -l) entries)"
fi

# ---- checkpoint reshape (stacked experts + bf16 PLE table) -------------------
# stack-qwen38-official-fp8.py reshapes the staged checkpoint in place
# (idempotent) and writes the PLE scale to $STAGE/ple_scale.txt — see the
# script's docstring for why each transform is value-preserving.
echo "stacking experts + merging the PLE table (value-preserving)..."
python3 "$SCRIPT_DIR/stack-qwen38-official-fp8.py" --stage "$STAGE"
PLE_SCALE="$(tr -d '\n' < "$STAGE/ple_scale.txt")"
echo "PLE fixed scale: $PLE_SCALE"

# ---- convert ----------------------------------------------------------------
# rad-convert needs no GPU: declare resolves kernels and asks each about the
# encoding it will read — CPU-side layout questions only. (If a plugin load
# ever demands a card, add --device /dev/kfd --device /dev/dri to this run.)
# --plan-only runs the same declare phase and exits before a byte is written.
echo "rad-convert: $SRC (staged) -> $OUT"
mkdir -p "$OUT_DIR"
CONVERT_ARGS=(/models/stage-qwen4exp-fp8 -o /models/rad/qwen4exp-fp8-native.rad
              --home /opt/radiance/share/radiance)
[ "$plan_only" -eq 1 ] && CONVERT_ARGS+=(--plan-only -v)
# THE ONE RECIPE RULE (see header): the PLE table arrives from staging as a
# bf16 merge of the checkpoint's own fp8 codes decoded with the checkpoint's
# own scale; rtn re-encodes it at that same scale — the codes come back
# exactly and the plugin's one-element scale plane is produced.
CONVERT_ARGS+=(--quant "blk.*.ple_ngram.weight=rtn:codes=fp8_e4m3,block=*x*,scale=bf16,rule=fixed,scale_value=$PLE_SCALE")
[ "$force" -eq 1 ] && [ -f "$OUT" ] && {
    TRASH="$HOME/.AGENTTRASH/$(date +%Y%m%d)-frank-archon-rad"
    mkdir -p "$TRASH"
    mv "$OUT" "$TRASH/qwen4exp-fp8-native.rad.$(date +%H%M%S)"
    echo "previous container moved to $TRASH"
}
podman run --rm \
    --entrypoint rad-convert \
    -v /home/dank/models:/models \
    "$IMAGE" "${CONVERT_ARGS[@]}"

if [ "$plan_only" -eq 1 ]; then
    echo "plan-only run complete (no container written)"
    exit 0
fi

# ---- record -----------------------------------------------------------------
echo "rad-info summary:"
podman run --rm --entrypoint rad-info -v /home/dank/models:/models:ro \
    "$IMAGE" "/models/rad/qwen4exp-fp8-native.rad" | head -20
MIB=$(( $(stat -c %s "$OUT") / 1048576 ))
echo "container: $OUT (${MIB} MiB; env floor $FRANK_ARCHON_RAD_MIN_MIB MiB)"
if [ "$MIB" -lt "$FRANK_ARCHON_RAD_MIN_MIB" ]; then
    echo "ERROR: container below the env floor — the convert was truncated" >&2
    exit 1
fi
echo "build complete"