#!/usr/bin/env bash
# build.sh - one-click build of stable-diffusion.cpp experiment variants for Z-Image-Turbo (CPU).
#
#   ./build.sh --all                  # deps + source + all variants + models
#   ./build.sh amx_lto avx512         # only the named variants
#   ./build.sh --list                 # show variants and their CMake flags
#   ./build.sh --skip-models amx_lto  # no model download
#   ./build.sh --skip-openvino --all  # skip the OpenVINO INT4 env + model (about 5.5 GB)
#
# Environment overrides:
#   WORK=./work                 working dir (source, binaries, models, venv)
#   SDCPP_REF=<commit>          upstream stable-diffusion.cpp commit (pinned by default)
#   JOBS=<n>                    parallel build jobs (default: nproc)
#   HF_TOKEN=<token>            only needed if a model repo becomes gated
#   VAE_REPO / VAE_FILE         VAE source (default: ffxvs/vae-flux/ae.safetensors, a mirror of the FLUX.1 VAE)
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORK="${WORK:-$PWD/work}"
SDCPP_REPO="${SDCPP_REPO:-https://github.com/leejet/stable-diffusion.cpp}"
SDCPP_REF="${SDCPP_REF:-77b72a2efed78f55de26892c7e909cee001e8b83}"
JOBS="${JOBS:-$(nproc)}"
VAE_REPO="${VAE_REPO:-ffxvs/vae-flux}"
VAE_FILE="${VAE_FILE:-ae.safetensors}"

SRC="$WORK/stable-diffusion.cpp"
BIN="$WORK/bin"
MODELS="$WORK/models"
VENV="$WORK/venv"

ALL_VARIANTS=(stock avx2 avx512 amx_lto llamafile)

variant_flags() {
    case "$1" in
        # upstream defaults: -march=native, no AMX, no LTO
        stock)   echo "-DGGML_NATIVE=ON" ;;
        # AVX2 only (AVX512 disabled), for measuring the AVX512 contribution
        avx2)    echo "-DGGML_NATIVE=OFF -DGGML_AVX=ON -DGGML_AVX2=ON -DGGML_FMA=ON -DGGML_F16C=ON -DGGML_BMI2=ON -DGGML_AVX512=OFF -DGGML_LTO=ON" ;;
        # explicit AVX512 + VNNI + BF16 + AMX (no -march=native)
        avx512)  echo "-DGGML_NATIVE=OFF -DGGML_AVX=ON -DGGML_AVX2=ON -DGGML_FMA=ON -DGGML_F16C=ON -DGGML_BMI2=ON -DGGML_AVX512=ON -DGGML_AVX512_VBMI=ON -DGGML_AVX512_VNNI=ON -DGGML_AVX512_BF16=ON -DGGML_AVX_VNNI=ON -DGGML_AMX_TILE=ON -DGGML_AMX_INT8=ON -DGGML_AMX_BF16=ON -DGGML_LTO=ON" ;;
        # native + AMX + LTO (best build on the reference machine)
        amx_lto) echo "-DGGML_NATIVE=ON -DGGML_AMX_TILE=ON -DGGML_AMX_INT8=ON -DGGML_AMX_BF16=ON -DGGML_LTO=ON" ;;
        # amx_lto + llamafile sgemm for F32 matmul/conv (VAE im2col path): stage 1
        llamafile) echo "-DGGML_NATIVE=ON -DGGML_AMX_TILE=ON -DGGML_AMX_INT8=ON -DGGML_AMX_BF16=ON -DGGML_LTO=ON -DGGML_LLAMAFILE=ON" ;;
        *) return 1 ;;
    esac
}

log()  { printf '\033[1;34m[build]\033[0m %s\n' "$*"; }
die()  { printf '\033[1;31m[error]\033[0m %s\n' "$*" >&2; exit 1; }

usage() {
    sed -n '2,14p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
    exit 0
}

# ---------------------------------------------------------------- args
DO_DEPS=1; DO_MODELS=1; DO_OV=1; SELECTED=()
for arg in "$@"; do
    case "$arg" in
        -h|--help) usage ;;
        --list)
            for v in "${ALL_VARIANTS[@]}"; do printf '%-9s %s\n' "$v" "$(variant_flags "$v")"; done
            exit 0 ;;
        --all) SELECTED=("${ALL_VARIANTS[@]}") ;;
        --skip-deps) DO_DEPS=0 ;;
        --skip-models) DO_MODELS=0 ;;
        --skip-openvino) DO_OV=0 ;;
        -*) die "unknown option: $arg" ;;
        *) variant_flags "$arg" >/dev/null || die "unknown variant: $arg (see --list)"; SELECTED+=("$arg") ;;
    esac
done
if [[ ${#SELECTED[@]} -eq 0 && $DO_MODELS -eq 1 ]]; then
    log "no variant given; use --all or name variants (see --list). Only fetching models."
fi

[[ "$(uname -m)" == "x86_64" ]] || log "warning: variants are tuned for x86_64 (AVX/AMX flags); this host is $(uname -m)"

mkdir -p "$WORK" "$BIN" "$MODELS"

# ---------------------------------------------------------------- deps
if [[ $DO_DEPS -eq 1 ]]; then
    if command -v apt-get >/dev/null 2>&1; then
        log "installing system packages (apt)"
        export DEBIAN_FRONTEND=noninteractive
        apt-get update -qq
        apt-get install -y -qq git cmake build-essential python3 python3-venv ca-certificates curl
    else
        log "apt-get not found; assuming git, cmake, a C++17 compiler and python3 (>=3.9) are installed"
    fi
    command -v git >/dev/null && command -v cmake >/dev/null || die "git and cmake are required"
    if [[ ! -x "$VENV/bin/python" ]]; then
        log "creating python venv at $VENV"
        python3 -m venv "$VENV"
    fi
    "$VENV/bin/pip" install -q --upgrade pip
    "$VENV/bin/pip" install -q huggingface_hub gguf numpy pillow
fi
PY="$VENV/bin/python"
[[ -x "$PY" ]] || PY="$(command -v python3)"

OV_VENV="$WORK/venv-ov"
OV_MODEL_REPO="${OV_MODEL_REPO:-HelloSun/Z-Image-Turbo-OpenVINO-INT4}"
OV_MODEL_DIR="$MODELS/Z-Image-Turbo-OpenVINO-INT4"
if [[ $DO_DEPS -eq 1 && $DO_OV -eq 1 ]]; then
    if [[ ! -x "$OV_VENV/bin/python" ]]; then
        log "creating OpenVINO venv at $OV_VENV (pinned: diffusers 0.37.1, optimum-intel 2.2.0, openvino 2026.4.0)"
        python3 -m venv "$OV_VENV"
        "$OV_VENV/bin/pip" install -q --upgrade pip
        "$OV_VENV/bin/pip" install -q 'diffusers==0.37.1' 'transformers==4.57.6' 'tokenizers==0.22.0' \
            'optimum==2.3.0' 'optimum-intel==2.2.0' 'openvino==2026.4.0' 'huggingface-hub==0.36.2' pillow numpy
    fi
fi

# ---------------------------------------------------------------- source
if [[ ${#SELECTED[@]} -gt 0 ]]; then
    if [[ ! -d "$SRC/.git" ]]; then
        log "cloning $SDCPP_REPO"
        git clone -q "$SDCPP_REPO" "$SRC"
    fi
    log "checking out sd.cpp $SDCPP_REF"
    git -C "$SRC" fetch -q origin || true
    git -C "$SRC" checkout -q -f "$SDCPP_REF"
    git -C "$SRC" reset -q --hard "$SDCPP_REF"
    git -C "$SRC" submodule update -q --init --recursive
    # discard leftovers of earlier patch runs so the patches below always apply to the pinned sources
    git -C "$SRC" submodule foreach -q --recursive 'git checkout -q -- . 2>/dev/null || true'

    # Apply our patches (idempotent: skip a patch that is already applied).
    apply_patch() {
        local repo="$1" patch="$2"
        if git -C "$repo" apply --reverse --check "$patch" >/dev/null 2>&1; then
            log "patch already applied: $(basename "$patch")"
        else
            git -C "$repo" apply --check "$patch" || die "patch does not apply: $(basename "$patch")"
            git -C "$repo" apply "$patch"
            log "applied patch: $(basename "$patch")"
        fi
    }
    apply_patch "$SRC"        "$ROOT/patches/sdcpp-0001-zimage-ffn-w13-and-cpu-extra-buft.patch"
    apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0001-cpu-op-profile.patch"
    apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0002-flash-attn-kpack.patch"
    apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0003-flash-attn-tiles.patch"
    apply_patch "$SRC/ggml"   "$ROOT/patches/ggml-0004-conv2d-im2col.patch"

    # ------------------------------------------------------------ build
    for v in "${SELECTED[@]}"; do
        log "building variant '$v': $(variant_flags "$v")"
        build_dir="$SRC/build-$v"
        # shellcheck disable=SC2086
        cmake -S "$SRC" -B "$build_dir" -DCMAKE_BUILD_TYPE=Release $(variant_flags "$v") >"$WORK/cmake-$v.log" 2>&1 \
            || { tail -n 30 "$WORK/cmake-$v.log"; die "cmake failed for $v"; }
        cmake --build "$build_dir" --config Release --target sd-cli -j "$JOBS" >"$WORK/build-$v.log" 2>&1 \
            || { tail -n 30 "$WORK/build-$v.log"; die "build failed for $v"; }
        mkdir -p "$BIN/$v"
        cp "$build_dir/bin/sd-cli" "$BIN/$v/sd-cli"
        log "-> $BIN/$v/sd-cli"
    done

    # q4k helper: wraps ggml's own Q4_K quantizer for tools/requant_refiner.py
    ref_build="$SRC/build-llamafile"
    [[ -f "$ref_build/ggml/src/libggml-base.a" ]] || ref_build="$SRC/build-${SELECTED[0]}"
    if [[ -f "$ref_build/ggml/src/libggml-base.a" ]]; then
        log "building q4k quantizer helper"
        gcc -O2 -fopenmp -I "$SRC/include" -I "$SRC/ggml/include" -I "$SRC/ggml/src" \
            "$ROOT/tools/q4k.c" "$ref_build/ggml/src/libggml-base.a" \
            -o "$WORK/q4k" -lpthread -lm -lstdc++ \
            || die "failed to build q4k helper"
    fi
fi

# ---------------------------------------------------------------- models
if [[ $DO_MODELS -eq 1 ]]; then
    log "downloading models into $MODELS"
    "$PY" - "$MODELS" "$VAE_REPO" "$VAE_FILE" <<'PY'
import os, sys
from huggingface_hub import hf_hub_download

models_dir, vae_repo, vae_file = sys.argv[1:4]
files = [
    ("leejet/Z-Image-Turbo-GGUF", "z_image_turbo-Q4_K.gguf"),
    ("leejet/Z-Image-Turbo-GGUF", "z_image_turbo-Q8_0.gguf"),
    ("leejet/Z-Image-Turbo-GGUF", "z_image_turbo-Q4_0.gguf"),
    ("unsloth/Qwen3-4B-Instruct-2507-GGUF", "Qwen3-4B-Instruct-2507-Q4_K_M.gguf"),
    (vae_repo, vae_file),
]
for repo, name in files:
    dst = os.path.join(models_dir, name)
    if os.path.exists(dst):
        print(f"  have {name}")
        continue
    print(f"  fetching {repo}/{name}")
    path = hf_hub_download(repo_id=repo, filename=name, local_dir=models_dir)
    print(f"  -> {path}")
PY
    if [[ $DO_OV -eq 1 ]]; then
        if [[ ! -f "$OV_MODEL_DIR/transformer/openvino_model.bin" ]]; then
            log "downloading OpenVINO INT4 model $OV_MODEL_REPO (about 5.5 GB)"
            "$PY" - "$OV_MODEL_REPO" "$OV_MODEL_DIR" <<'PY'
import sys
from huggingface_hub import snapshot_download
repo, dst = sys.argv[1:3]
snapshot_download(repo_id=repo, local_dir=dst,
                  ignore_patterns=["examples/*", "generated/*", "*.png", "README.md", "REPORT.md"])
print("  ->", dst)
PY
        else
            log "have OpenVINO model at $OV_MODEL_DIR"
        fi
    fi
    # FFN gate/up fusion at the weight level (w1 + w3 -> w13); see tools/fuse_ffn.py
    if [[ ! -f "$MODELS/z_image_turbo-Q4_K-ffn13.gguf" ]]; then
        log "fusing FFN w1/w3 -> w13 (Q4_K)"
        "$PY" "$ROOT/tools/fuse_ffn.py" "$MODELS/z_image_turbo-Q4_K.gguf" "$MODELS/z_image_turbo-Q4_K-ffn13.gguf"
    fi
    # refiner requant: the reference GGUF stores the 2 noise + 2 context refiner
    # blocks as Q8_0, which is ~2x slower than Q4_K on a memory-bound CPU GEMM.
    # Rewrite just those 16 attention/FFN matrices as Q4_K (adaLN left at Q8_0).
    if [[ -x "$WORK/q4k" && ! -f "$MODELS/z_image_turbo-Q4_K-ffn13-rq.gguf" ]]; then
        log "requantizing refiner blocks Q8_0 -> Q4_K"
        Q4K="$WORK/q4k" "$PY" "$ROOT/tools/requant_refiner.py" \
            "$MODELS/z_image_turbo-Q4_K-ffn13.gguf" "$MODELS/z_image_turbo-Q4_K-ffn13-rq.gguf"
    fi
fi

log "done"
for v in "${ALL_VARIANTS[@]}"; do
    [[ -x "$BIN/$v/sd-cli" ]] && log "binary: $BIN/$v/sd-cli"
done
exit 0