#!/bin/sh
# usage: sh install-comfyui-depends-qwen-image-21.sh [--force] [--delete] [comfyui_root]
#   --force      refetch everything, ignore size match, never prompt
#   --delete     remove the model files and exit, no network access
#   ASSUME_YES=1 answer yes to mismatch prompts when not on a terminal
#   HF_TOKEN=hf_...   bearer token, forwarded across the CDN redirect
#   DEL=1         alias for --force
set -eu

BASE="https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/main"

FORCE=0
MODE=fetch
ROOT="$HOME/AI/ComfyUI"
for arg in "$@"; do
    case "$arg" in
        -f|--force) FORCE=1 ;;
        -d|--delete) MODE=delete ;;
        -*) echo "unknown option: $arg" >&2; exit 2 ;;
        *) ROOT="$arg" ;;
    esac
done
if [ "${DEL:-0}" = 1 ]; then
    FORCE=1
fi

cd "$ROOT"

N_NEW=0
N_SKIP=0
N_UPD=0
N_BAD=0
N_DEL=0

human() {
    awk -v b="$1" 'BEGIN {
        split("B KiB MiB GiB TiB", u, " ")
        i = 1
        while (b >= 1024 && i < 5) { b /= 1024; i++ }
        printf (i == 1 ? "%d %s" : "%.1f %s"), b, u[i]
    }'
}

fsize() {
    wc -c < "$1" | awk '{ print $1 }'
}

name_of() {
    name=${1%%\?*}
    name=${name%%#*}
    echo "${name##*/}"
}

want_size() {
    curl -sIL "$1" | tr -d '\r' |
        awk 'tolower($1) == "x-linked-size:" { v = $2 } END { print v }'
}

confirm() {
    if [ ! -t 0 ] && [ "${ASSUME_YES:-0}" != 1 ]; then
        echo "   (no terminal, ASSUME_YES!=1, not touching it)" >&2
        return 1
    fi
    if [ -t 0 ]; then
        printf '%s [y/N] ' "$1"
        ans=""
        read -r ans || ans=""
        case "$ans" in
            y|Y|yes|YES|Yes) return 0 ;;
            *) return 1 ;;
        esac
    fi
    return 0
}

remove() {
    path="$2/$(name_of "$1")"
    part="$path.part"
    echo "== $path"
    if [ -f "$path" ]; then
        echo "   deleting $path ($(human "$(fsize "$path")"))"
        rm -f "$path"
        N_DEL=$((N_DEL + 1))
    else
        echo "   nothing to delete"
    fi
    if [ -f "$part" ]; then
        echo "   deleting $part"
        rm -f "$part"
    fi
    return 0
}

fetch() {
    url=$1
    dir=$2
    path="$dir/$(name_of "$url")"
    part="$path.part"

    set -- -fL --progress-bar --retry 5 --retry-delay 3 --retry-all-errors \
        --connect-timeout 30 --create-dirs
    if [ -n "${HF_TOKEN:-}" ]; then
        set -- "$@" --location-trusted -H "Authorization: Bearer $HF_TOKEN"
    fi

    echo "== $path"
    echo "   url:    $url"

    want=$(want_size "$url" || true)
    if [ -n "$want" ]; then
        echo "   expect: $want bytes ($(human "$want"))"
    else
        echo "   expect: unknown, server published no size"
    fi

    action=install
    if [ -f "$path" ]; then
        have=$(fsize "$path")
        echo "   have:   $have bytes ($(human "$have"))"
        if [ "$FORCE" = 1 ]; then
            echo "   --force, removing and refetching" >&2
            rm -f "$path"
            action=update
        elif [ -z "$want" ]; then
            echo "   warn:   no size to compare against, cannot verify"
            N_SKIP=$((N_SKIP + 1))
            return 0
        elif [ "$have" = "$want" ]; then
            echo "   ok:     size matches, skipping download"
            N_SKIP=$((N_SKIP + 1))
            return 0
        elif confirm "   size mismatch. Remove and refetch $(human "$want")?"; then
            echo "   BAD:    size mismatch, removing and refetching" >&2
            rm -f "$path"
            action=update
        else
            echo "   keep:   left $path unchanged" >&2
            N_BAD=$((N_BAD + 1))
            return 1
        fi
    fi

    echo "   fetching to $part"
    rm -f "$part"
    set -- "$@" -o "$part"
    if ! curl "$@" "$url"; then
        echo "   error:  download failed, discarded $part" >&2
        rm -f "$part"
        N_BAD=$((N_BAD + 1))
        return 1
    fi

    have=$(fsize "$part")
    echo "   have:   $have bytes ($(human "$have"))"
    if [ -n "$want" ] && [ "$have" != "$want" ]; then
        echo "   BAD:    size mismatch, discarded $part" >&2
        rm -f "$part"
        N_BAD=$((N_BAD + 1))
        return 1
    fi

    mv "$part" "$path"
    if [ "$action" = update ]; then
        echo "   done:   updated $path"
        N_UPD=$((N_UPD + 1))
    else
        echo "   done:   installed $path"
        N_NEW=$((N_NEW + 1))
    fi
    return 0
}

all() {
    if [ "$MODE" = delete ]; then
        remove "$@"
    else
        fetch "$@" || return 1
    fi
}

FAILED=0
all "$BASE/vae/qwen_image_2.1_vae_bf16.safetensors" models/vae || FAILED=1
all "$BASE/text_encoders/qwen3vl_8b_int8_convrot.safetensors" models/text_encoders || FAILED=1
all "$BASE/text_encoders/qwen3.5_9b_qwen_image_2.1_pe_t2i.int8_convrot.safetensors" models/text_encoders || FAILED=1
all "$BASE/text_encoders/qwen3.5_9b_qwen_image_2.1_pe_i2i.int8_convrot.safetensors" models/text_encoders || FAILED=1
all "$BASE/diffusion_models/qwen_image_2.1_int8_convrot.safetensors" models/diffusion_models || FAILED=1

echo
if [ "$MODE" = delete ]; then
    echo "summary: $N_DEL deleted"
    exit 0
fi
echo "summary: $N_NEW installed, $N_UPD updated, $N_SKIP already present, $N_BAD bad"
[ "$FAILED" = 0 ] || exit 1