#!/usr/bin/env bash
set -euo pipefail
# POS: ai hf — Download AI models from Hugging Face (search, download, manage)
# POS_SUBCMDS: search download list remove info files cache
# POS_FLAGS: --branch --gguf --list --output --quant --include --exclude --revision
# POS_DEPS: curl jq
# POS_CONFIG: ai | ai.env | HF_TOKEN=secret:Hugging Face API token (https://huggingface.co/settings/tokens) | HF_DOWNLOAD_DIR=:Model download directory (default ~/.local/share/linux_post_install/ai/models)
# POS_EXAMPLES: pos ai hf search llama 7b | Search Hugging Face for "llama 7b" models
# POS_EXAMPLES: pos ai hf download meta-llama/Llama-3.1-8B-Instruct | Download all files from a repo
# POS_EXAMPLES: pos ai hf download meta-llama/Llama-3.1-8B-Instruct --gguf | Download only GGUF quantized files
# POS_EXAMPLES: pos ai hf download org/model-GGUF --gguf --quant Q8_0 | Download one quant directory's GGUF shards
# POS_EXAMPLES: pos ai hf download meta-llama/Llama-3.1-8B-Instruct --list | List remote repository files (what --gguf/download would fetch)
# POS_EXAMPLES: pos ai hf download meta-llama/Llama-3.1-8B-Instruct config.json | Download a single file
# POS_EXAMPLES: pos ai hf list | List downloaded models
# POS_EXAMPLES: pos ai hf remove meta-llama-Llama-3.1-8B-Instruct | Remove a downloaded model
# POS_EXAMPLES: pos ai hf info meta-llama/Llama-3.1-8B-Instruct | Show repository information
# POS_EXAMPLES: pos ai hf files meta-llama/Llama-3.1-8B-Instruct | List repository files
# POS_EXAMPLES: pos ai hf download meta-llama/Llama-3.1-8B-Instruct --include "*.gguf" --exclude "*Q4_*" | Download with include/exclude patterns

source "$(dirname "$0")/../lib/common.sh" 2>/dev/null || source "$(dirname "$0")/common.sh"

# Shared config loader (canonical KEY=VALUE parser, env-wins precedence)
source "$(dirname "$0")/../lib/config-ui.sh" 2>/dev/null || source "$(dirname "$0")/config-ui.sh"

# ── Dependencies (before --help) ───────────────────────────────
command -v curl &>/dev/null || err "curl not found (install curl)"
command -v jq &>/dev/null   || err "jq not found (install jq)"

# ── Config & defaults ──────────────────────────────────────────
CONFIG_FILE="${CONFIG_FILE:-$CONFIG_DIR/ai.env}"
HF_TOKEN="${HF_TOKEN:-}"
HF_DOWNLOAD_DIR="${HF_DOWNLOAD_DIR:-$HOME/.local/share/linux_post_install/ai/models}"

load_hf_config() {
    load_env_file "$CONFIG_FILE"
}

load_hf_config

# ── Usage ──────────────────────────────────────────────────────
usage() {
    cat <<'EOF'
Usage: pos ai hf <subcommand> [args]

Hugging Face model downloader — search, download, and manage AI models.

Subcommands:
  search <query>                     Search Hugging Face models
  download <repo-id> [filename]      Download a file or entire repo
  list                               List locally downloaded models
  remove <repo-id>                   Remove a downloaded model
  info <repo-id>                     Show repository information
  files <repo-id>                    List repository files
  cache [status|clear]               Inspect cache (status) or remove all
                                     downloaded models (clear — asks for
                                     confirmation)

Download options:
  --branch <rev>                     Download from a specific branch/revision
  --gguf                             Download only .gguf weight files (excludes
                                     mmproj/imatrix/vision/MTP artifacts)
  --quant <dir>                      With --gguf: pick one quant directory when a
                                     repo groups weights into several (e.g.
                                     --gguf --quant Q8_0)
  --list                             List remote repository files without downloading
  --output <dir>                     Override download directory
  --include <pattern>                Include files matching pattern (supports glob)
  --exclude <pattern>                Exclude files matching pattern (supports glob)
  --revision <rev>                   Specific revision (commit/tag/branch);
                                     alias for --branch — when both are given,
                                     the later one wins

Examples:
  pos ai hf search llama 7b
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct --gguf
  pos ai hf download org/model-GGUF --gguf --quant Q8_0
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct --list
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct config.json
  pos ai hf download org/model-GGUF Q8_0/model-00001-of-00006.gguf
  pos ai hf download org/model-GGUF model-00001-of-00006.gguf
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct --branch main
  pos ai hf list
  pos ai hf remove meta-llama-Llama-3.1-8B-Instruct
  pos ai hf info meta-llama/Llama-3.1-8B-Instruct
  pos ai hf files meta-llama/Llama-3.1-8B-Instruct
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct --include "*.gguf" --exclude "*Q4_*"
  pos ai hf download meta-llama/Llama-3.1-8B-Instruct --revision v1.0

A filename may be a full path (Q8_0/model.gguf) or a bare name (model.gguf) —
bare names matching files in multiple directories error and ask for the full path.
--list shows files on the remote repo; 'list' shows models already downloaded.

Config (~/.config/linux_post_install/ai.env):
  HF_TOKEN           Hugging Face API token (better rate limits for public repos)
  HF_DOWNLOAD_DIR    Model download directory (default ~/.local/share/linux_post_install/ai/models)

Exit codes:
  0  success
  1  error (missing deps, invalid input, API failure)
EOF
    exit 0
}

# ── Parse global flags ─────────────────────────────────────────
SUBCMD=""
SUBCMD_ARGS=()
GGUF_ONLY=0
OUTPUT_DIR=""
LIST_FILES=0
QUANT_DIR=""
INCLUDE_PATTERN=""
EXCLUDE_PATTERN=""
REVISION=""

while [ $# -gt 0 ]; do
    case "$1" in
        -h|--help) usage ;;
        --branch)
            [ $# -ge 2 ] || err "--branch requires a value"
            REVISION="$2"; shift 2 ;;
        --gguf)
            GGUF_ONLY=1; shift ;;
        --list)
            LIST_FILES=1; shift ;;
        --quant)
            [ $# -ge 2 ] || err "--quant requires a value"
            QUANT_DIR="$2"; shift 2 ;;
        --output)
            [ $# -ge 2 ] || err "--output requires a value"
            OUTPUT_DIR="$2"; shift 2 ;;
        --include)
            [ $# -ge 2 ] || err "--include requires a value"
            INCLUDE_PATTERN="$2"; shift 2 ;;
        --exclude)
            [ $# -ge 2 ] || err "--exclude requires a value"
            EXCLUDE_PATTERN="$2"; shift 2 ;;
        --revision)
            [ $# -ge 2 ] || err "--revision requires a value"
            REVISION="$2"; shift 2 ;;
        -*)
            err "Unknown option '$1' (see --help)" ;;
        *)
            if [ -z "$SUBCMD" ]; then
                SUBCMD="$1"
            else
                SUBCMD_ARGS+=("$1")
            fi
            shift ;;
    esac
done

[ -n "$SUBCMD" ] || usage

# Apply output dir override
if [ -n "$OUTPUT_DIR" ]; then
    HF_DOWNLOAD_DIR="$OUTPUT_DIR"
fi
HF_DOWNLOAD_DIR="${HF_DOWNLOAD_DIR:-$HOME/.local/share/linux_post_install/ai/models}"

# ── Token warning ──────────────────────────────────────────────
if [ -z "$HF_TOKEN" ]; then
    warn "No HF_TOKEN set — using anonymous access"
fi

# ── HF API helpers ─────────────────────────────────────────────
HF_BASE="https://huggingface.co"
HF_API_BASE="https://huggingface.co/api"

HF_MAX_PAGES=20
HF_GGUF_FILTER='[ .[] |
  select(.rfilename | type == "string") |
  select(.rfilename | ascii_downcase | endswith(".gguf")) |
  select(.rfilename | ascii_downcase | test("mmproj|imatrix|clip|vision|projector|mtp") | not)
]'

hf_auth_header() {
    if [ -n "$HF_TOKEN" ]; then
        printf 'Authorization: Bearer %s' "$HF_TOKEN"
    fi
}

hf_api() {
    local endpoint="$1"
    local hdr_file="${2:-}"          # optional: dump response headers (Link: rel="next")
    local url
    case "$endpoint" in
        http://*|https://*) url="$endpoint" ;;
        *)                  url="${HF_API_BASE}${endpoint}" ;;
    esac
    local auth_header
    auth_header="$(hf_auth_header)"

    local http_code body tmpfile
    tmpfile="$(mktemp)"

    local curl_args=(-sS -w '%{http_code}' -o "$tmpfile" --max-time 30)
    if [ -n "$auth_header" ]; then
        curl_args+=(-H "$auth_header")
    fi
    if [ -n "$hdr_file" ]; then
        curl_args+=(-D "$hdr_file")
    fi

    # Rate limit retry: on 429, sleep and retry once
    local attempt=0
    while [ $attempt -lt 2 ]; do
        http_code="$(curl "${curl_args[@]}" "$url" 2>/dev/null)" || {
            rm -f "$tmpfile" "$hdr_file"
            err "Connection timed out — check network"
        }

        if [ "$http_code" = "429" ]; then
            local retry_after
            retry_after="$(curl -sI -H "${auth_header:-}" "$url" 2>/dev/null | grep -i 'retry-after:' | tr -d '\r' | awk '{print $2}')"
            retry_after="${retry_after:-60}"
            warn "Rate limited — waiting ${retry_after}s before retry"
            sleep "$retry_after"
            attempt=$((attempt + 1))
            continue
        fi

        break
    done

    body="$(cat "$tmpfile")"
    rm -f "$tmpfile"

    case "$http_code" in
        200) ;;
        401|403) err "Authentication failed — check HF_TOKEN (pos config ai)" ;;
        404) err "Model not found: ${endpoint#/api/models/}" ;;
        429) err "Rate limit exceeded — try again later" ;;
        *)   err "API request failed (HTTP $http_code)" ;;
    esac

    # Validate JSON
    if ! printf '%s' "$body" | jq empty 2>/dev/null; then
        err "Failed to parse API response — check network or HF status"
    fi

    printf '%s' "$body"
}

# hf_paginate <endpoint> → JSON array built from every Link: rel="next" page
hf_paginate() {
    local endpoint="$1"
    local url
    case "$endpoint" in
        http://*|https://*) url="$endpoint" ;;
        *)                  url="${HF_API_BASE}${endpoint}" ;;
    esac
    local combined="[]"
    local page=0
    local hdr_file body next_url
    while [ -n "$url" ]; do
        page=$((page + 1))
        [ "$page" -gt "$HF_MAX_PAGES" ] \
            && err "Repository listing exceeded ${HF_MAX_PAGES} pages — aborting"
        hdr_file="$(mktemp)"
        body="$(hf_api "$url" "$hdr_file")"
        combined="$(printf '%s\n%s' "$combined" "$body" | jq -c -s 'add')"
        next_url="$(sed -n 's/^link: <\([^>]*\)>; rel="next".*/\1/Ip' "$hdr_file" | tr -d '\r' | tail -1)"
        rm -f "$hdr_file"
        url="${next_url:-}"
    done
    printf '%s' "$combined"
}

hf_repo_files() {
    local repo_id="$1"
    local branch="${2:-main}"
    local ns="${repo_id%%/*}"
    local repo="${repo_id#*/}"

    if [ "$ns" = "$repo" ]; then
        err "Invalid repo format: use namespace/model-name"
    fi

    # Try /tree/ endpoint first (has file sizes + LFS info)
    local endpoint="/models/${ns}/${repo}/tree/${branch}?recursive=true"
    local result
    if result="$(hf_paginate "$endpoint" 2>/dev/null)"; then
        # Tree API returns {type,path,size,oid[,lfs]} per entry — normalize to the
        # {rfilename,size} shape the rest of the pipeline expects (same as fallback).
        # Skip "directory" entries and guard non-object entries (error objects crash .[]).
        printf '%s' "$result" | jq '[.[] | select(type == "object" and .type == "file") | {rfilename: .path, size: (.size // 0)}]'
        return 0
    fi

    # Fallback: /api/models/{ns}/{repo} (siblings, no sizes)
    warn "Tree endpoint unavailable, using repo metadata"
    local fallback
    fallback="$(hf_api "/models/${ns}/${repo}")" || err "Failed to fetch repo info for $repo_id"
    printf '%s' "$fallback" | jq '[.siblings[]? | select(type == "object") | {rfilename: (.rfilename // ""), size: (.size // 0)}]'
}

hf_search() {
    local query="$1"
    local limit="${2:-10}"
    local encoded_query
    encoded_query="$(printf '%s' "$query" | jq -sRr @uri)"
    local endpoint="/models?search=${encoded_query}&sort=downloads&direction=-1&limit=${limit}"
    hf_api "$endpoint"
}

# ── Utility functions ──────────────────────────────────────────
hf_repo_dir() {
    local repo_id="$1"
    printf '%s' "$HF_DOWNLOAD_DIR/${repo_id//\//-}"
}

hf_human_size() {
    local bytes="$1"
    if [ "$bytes" -ge 1073741824 ]; then
        awk "BEGIN { printf \"%.1f GB\", $bytes / 1073741824 }"
    elif [ "$bytes" -ge 1048576 ]; then
        awk "BEGIN { printf \"%.1f MB\", $bytes / 1048576 }"
    elif [ "$bytes" -ge 1024 ]; then
        awk "BEGIN { printf \"%.1f KB\", $bytes / 1024 }"
    else
        printf '%d B' "$bytes"
    fi
}

# hf_quant_candidates <files-json> → [{dir, files, size}] sorted by dir
hf_quant_candidates() {
    printf '%s' "$1" | jq -c '[.[] | select(.rfilename | contains("/")) |
      {dir: (.rfilename | split("/")[0]), size: (.size // 0)}]
      | group_by(.dir)
      | map({dir: .[0].dir, files: length, size: (map(.size) | add)})
      | sort_by(.dir)'
}

# Refactored hf_gguf_quant_gate function with improved structure
# hf_gguf_quant_gate <files-json> <quant-dir> <repo-id> → filtered JSON (stdout) or err
hf_gguf_quant_gate() {
    local json="$1" quant="${2:-}" repo_id="$3"
    
    # Validate input
    if [ -z "$json" ]; then
        err "No files provided to quant gate"
    fi
    
    # Count top-level files vs directory files
    local top_count dir_count
    top_count="$(printf '%s' "$json" | jq '[.[] | select(.rfilename | contains("/") | not)] | length')"
    dir_count="$(printf '%s' "$json" | jq '[.[] | select(.rfilename | contains("/")) | .rfilename | split("/")[0]] | unique | length')"

    # Handle case: top-level .gguf files (no quant dirs)
    if [ "$top_count" -gt 0 ]; then
        if [ -n "$quant" ]; then
            err "--quant is for repos that group weights into quant directories — $repo_id has top-level .gguf files, --quant is not needed"
        fi
        printf '%s' "$json"
        return 0
    fi

    # Handle case: single quant directory
    if [ "$dir_count" -eq 1 ]; then
        local only_dir
        only_dir="$(printf '%s' "$json" | jq -r '.[0].rfilename | split("/")[0]')"
        if [ -n "$quant" ] && [ "$quant" != "$only_dir" ]; then
            err "No quant directory '$quant' in $repo_id — weights live in: $only_dir"
        fi
        printf '%s' "$json"
        return 0
    fi

    # Handle case: multiple quant directories - require quant selection
    if [ -z "$quant" ]; then
        local msg
        msg="$(printf 'Repo %s organizes weights into %d quant directories — pick one with --quant:\n' "$repo_id" "$dir_count")"
        while IFS=$'\t' read -r dir files size; do
            msg+="$(printf '  %-20s %d files, %s\n' "$dir" "$files" "$(hf_human_size "$size")")"
        done < <(hf_quant_candidates "$json" | jq -r '.[] | [.dir, (.files|tostring), (.size|tostring)] | @tsv')
        err "$msg"
    fi

    # Filter by specified quant directory
    local selected
    selected="$(printf '%s' "$json" | jq -c --arg q "$quant" '[.[] | select(.rfilename | split("/")[0] == $q)]')"
    if [ "$(printf '%s' "$selected" | jq 'length')" -eq 0 ]; then
        local msg
        msg="$(printf 'No weights in quant directory %s in %s — candidates:\n' "$quant" "$repo_id")"
        while IFS=$'\t' read -r dir files size; do
            msg+="$(printf '  %-20s %d files, %s\n' "$dir" "$files" "$(hf_human_size "$size")")"
        done < <(hf_quant_candidates "$json" | jq -r '.[] | [.dir, (.files|tostring), (.size|tostring)] | @tsv')
        err "$msg"
    fi
    printf '%s' "$selected"
}

# hf_list_files <repo-id> <branch> <files-json> → stdout table, no downloads
hf_list_files() {
    local repo_id="$1" branch="$2" json="$3"
    local count total
    count="$(printf '%s' "$json" | jq 'length')"
    [ "$count" -gt 0 ] || err "No files found in $repo_id${branch:+ (branch: $branch)}"
    total="$(printf '%s' "$json" | jq '[.[].size // 0] | add // 0')"
    printf 'Files in %s (branch: %s, %d file(s), %s):\n' \
        "$repo_id" "$branch" "$count" "$(hf_human_size "$total")"
    printf '%s' "$json" | jq -r 'sort_by(.rfilename)[] | [.rfilename, (.size // 0)] | @tsv' | \
        while IFS=$'\t' read -r rpath rsize; do
            printf '  %-60s %s\n' "$rpath" "$(hf_human_size "$rsize")"
        done
}

# hf_apply_patterns <files-json> <include-pattern> <exclude-pattern> → filtered
# JSON array. Glob semantics via bash `case` (the documented "supports glob"):
# keep entries whose rfilename matches $include (when set) AND does not match
# $exclude (when set). Composes after the gguf/filename filters and preserves
# the {"rfilename","size"} array shape downstream consumers expect.
hf_apply_patterns() {
    local json="$1" include="${2:-}" exclude="${3:-}"
    local entries=() entry fname
    while IFS= read -r entry; do
        [ -n "$entry" ] || continue
        fname="$(printf '%s' "$entry" | jq -r '.rfilename // empty')"
        [ -n "$fname" ] || continue
        if [ -n "$include" ]; then
            case "$fname" in
                $include) ;;
                *) continue ;;
            esac
        fi
        if [ -n "$exclude" ]; then
            case "$fname" in
                $exclude) continue ;;
            esac
        fi
        entries+=("$entry")
    done < <(printf '%s' "$json" | jq -c '.[]')
    if [ "${#entries[@]}" -eq 0 ]; then
        printf '[]'
    else
        printf '%s\n' "${entries[@]}" | jq -c -s '.'
    fi
}

hf_resolve_branch() {
    local repo_id="$1"
    local branch="${2:-}"
    if [ -n "$branch" ]; then
        printf '%s' "$branch"
        return
    fi
    # Try to get default branch from API
    local ns="${repo_id%%/*}"
    local repo="${repo_id#*/}"
    local meta
    if meta="$(hf_api "/models/${ns}/${repo}" 2>/dev/null)"; then
        local default_branch
        default_branch="$(printf '%s' "$meta" | jq -r '.defaultBranch // empty' 2>/dev/null)"
        if [ -n "$default_branch" ]; then
            printf '%s' "$default_branch"
            return
        fi
    fi
    printf 'main'
}

# ── Download helper ────────────────────────────────────────────
# Enhanced progress function to provide better feedback
hf_download_with_progress() {
    local url="$1"
    local target="$2"
    local file_name="$(basename "$target")"
    
    # Create parent directory
    mkdir -p "$(dirname "$target")"
    
    local auth_header
    auth_header="$(hf_auth_header)"
    
    local curl_args=(-L -C - --progress-bar -o "$target")
    if [ -n "$auth_header" ]; then
        curl_args+=(-H "$auth_header")
    fi
    
    # Run download with progress bar
    if curl "${curl_args[@]}" "$url" 2>&1; then
        if [ -s "$target" ]; then
            return 0
        else
            warn "Downloaded file is empty: $target"
            return 1
        fi
    else
        warn "Download interrupted for $file_name (resume with same command)"
        return 1
    fi
}

# ── Parallel download limit (used by cmd_download) ─────────────
PARALLEL_DOWNLOADS=4  # Default parallel downloads

# ── Subcommands ────────────────────────────────────────────────

cmd_search() {
    local query="${SUBCMD_ARGS[*]:-}"
    [ -n "$query" ] || err "Usage: pos ai hf search <query>"

    local result
    result="$(hf_search "$query" "10")"

    local count
    count="$(printf '%s' "$result" | jq 'length')"
    [ "$count" -gt 0 ] || { warn "No models found for '$query'"; return 0; }

    printf 'Found %d models for "%s":\n' "$count" "$query"
    printf '%s' "$result" | jq -r '.[] | "  \(.id)\t\(.downloads // 0)\t\(.likes // 0)"' | \
        while IFS=$'\t' read -r id downloads likes; do
            local dl_str
            if [ "$downloads" -ge 1000 ]; then
                dl_str="$(awk "BEGIN { printf \"%.1fk\", $downloads / 1000 }")"
            else
                dl_str="$downloads"
            fi
            printf '  %-55s %s downloads\n' "$id" "$dl_str"
        done
}

cmd_download() {
    local repo_id="${SUBCMD_ARGS[0]:-}"
    [ -n "$repo_id" ] || err "Usage: pos ai hf download <repo-id> [filename]"

    # Validate repo-id contains /
    [[ "$repo_id" == */* ]] || err "Invalid repo format: use namespace/model-name"

    local filename="${SUBCMD_ARGS[1]:-}"

    # Flag pre-checks
    [ -n "$QUANT_DIR" ] && [ "$GGUF_ONLY" -eq 0 ] && err "--quant requires --gguf"
    [ "$LIST_FILES" -eq 1 ] && [ -n "$filename" ] && err "--list cannot be combined with a filename"
    [ -n "$INCLUDE_PATTERN" ] && [ -n "$EXCLUDE_PATTERN" ] && [ -n "$filename" ] && err "--include/--exclude cannot be used with specific filenames"

    local branch
    branch="$(hf_resolve_branch "$repo_id" "$REVISION")"

    # Get file list from API (recursive + paginated tree)
    local files_json
    files_json="$(hf_repo_files "$repo_id" "$branch")"

    # --list mode: print what download would fetch, don't download
    if [ "$LIST_FILES" -eq 1 ]; then
        local list_json="$files_json"
        if [ "$GGUF_ONLY" -eq 1 ]; then
            list_json="$(printf '%s' "$list_json" | jq -c "$HF_GGUF_FILTER")"
            [ "$(printf '%s' "$list_json" | jq 'length')" -gt 0 ] \
                && list_json="$(hf_gguf_quant_gate "$list_json" "$QUANT_DIR" "$repo_id")"
        fi
        hf_list_files "$repo_id" "$branch" "$list_json"
        return 0
    fi

    # Filter files
    local filtered_files
    if [ -n "$filename" ]; then
        # Single file mode — explicit filename wins over --gguf/--quant
        if [[ "$filename" == */* ]]; then
            # Full path → exact .rfilename match
            filtered_files="$(printf '%s' "$files_json" | jq -c --arg fn "$filename" '[.[] | select(.rfilename == $fn)]')"
        else
            # Bare name → basename match across all depths
            filtered_files="$(printf '%s' "$files_json" | jq -c --arg fn "$filename" '[.[] | select((.rfilename | type) == "string") | select(.rfilename | split("/")[-1] == $fn)]')"
        fi
    elif [ "$GGUF_ONLY" -eq 1 ]; then
        # GGUF filter
        filtered_files="$(printf '%s' "$files_json" | jq -c "$HF_GGUF_FILTER")"
        [ "$(printf '%s' "$filtered_files" | jq 'length')" -gt 0 ] \
            && filtered_files="$(hf_gguf_quant_gate "$filtered_files" "$QUANT_DIR" "$repo_id")"
    else
        # All files
        filtered_files="$(printf '%s' "$files_json" | jq -c '.')"
    fi

    # Pattern filtering (bash `case` glob) — composes with the filters above in
    # the order gguf/filename → include → exclude, and always yields an array
    if [ -n "$INCLUDE_PATTERN" ] || [ -n "$EXCLUDE_PATTERN" ]; then
        filtered_files="$(hf_apply_patterns "$filtered_files" "$INCLUDE_PATTERN" "$EXCLUDE_PATTERN")"
    fi

    local file_count
    file_count="$(printf '%s' "$filtered_files" | jq 'length')"
    if [ "$file_count" -eq 0 ]; then
        if [ -n "$filename" ]; then
            err "File not found: $filename in $repo_id (branch: ${branch})"
        elif [ "$GGUF_ONLY" -eq 1 ]; then
            err "No .gguf files found in $repo_id${branch:+ (branch: $branch)} — try without --gguf"
        elif [ -n "$INCLUDE_PATTERN" ] || [ -n "$EXCLUDE_PATTERN" ]; then
            err "No files match include/exclude patterns in $repo_id${branch:+ (branch: $branch)}"
        else
            err "No files to download"
        fi
    fi

    # Ambiguity guard: bare name matching multiple files (subdirs) → ask for full path
    if [ -n "$filename" ] && [[ "$filename" != */* ]] && [ "$file_count" -gt 1 ]; then
        err "$(printf 'Multiple files match "%s" in %s — use the full path:\n' "$filename" "$repo_id"; printf '%s' "$filtered_files" | jq -r '.[] | "  \(.rfilename)"')"
    fi

    # Prepare target directory
    local target_dir
    target_dir="$(hf_repo_dir "$repo_id")"
    mkdir -p "$target_dir"

    # Disk space pre-flight check
    local total_size
    total_size="$(printf '%s' "$filtered_files" | jq '[.[].size // 0] | add // 0')"
    if [ "$total_size" -gt 0 ]; then
        local avail_kb
        avail_kb="$(df --output=avail "$target_dir" 2>/dev/null | tail -1 | tr -d ' ')"
        local need_kb=$((total_size / 1024))
        if [ "$avail_kb" -lt "$need_kb" ]; then
            local need_human avail_human
            need_human="$(hf_human_size "$total_size")"
            avail_human="$(hf_human_size "$((avail_kb * 1024))")"
            warn "Low disk space: need $need_human, only $avail_human available"
        fi
    fi

    local downloaded=0
    local failed_files=()
    local ns="${repo_id%%/*}"
    local repo="${repo_id#*/}"

    # Multiple files → parallel downloads. Each job is reaped individually so
    # one failure does not abort the batch; failures are collected in
    # failed_files (function-scoped so the summary/meta/exit below can see
    # them), reported per-file at the end, and the temp dir is removed on
    # every exit path (EXIT trap).
    if [ "$file_count" -gt 1 ]; then
        local temp_dir
        temp_dir="$(mktemp -d)"
        trap 'rm -rf "$temp_dir"' EXIT
        local job_pids=() job_names=()
        local max_jobs="${PARALLEL_DOWNLOADS:-4}"
        local completed_jobs=0 n=0 i

        # Process files in parallel batches
        while IFS= read -r file_json; do
            local fname fsize
            fname="$(printf '%s' "$file_json" | jq -r '.rfilename')"
            fsize="$(printf '%s' "$file_json" | jq -r '.size // 0' || echo 0)"

            local url="${HF_BASE}/${ns}/${repo}/resolve/${branch}/${fname}"
            local target="${target_dir}/${fname}"

            # Start background job; each job logs to its own file so output
            # does not interleave
            ( hf_download_with_progress "$url" "$target" >"$temp_dir/job-$n.log" 2>&1 ) &
            local pid=$!
            job_pids+=("$pid")
            job_names+=("$fname")
            n=$((n + 1))

            # Limit parallel jobs — reap the oldest job; a failed download is
            # recorded and labelled honestly, never fatal to the batch
            if [ "${#job_pids[@]}" -ge "$max_jobs" ]; then
                completed_jobs=$((completed_jobs + 1))
                if ! wait "${job_pids[0]}"; then
                    failed_files+=("${job_names[0]}")
                    printf '[%d/%d] Failed: %s\n' "$completed_jobs" "$file_count" "${job_names[0]}" >&2
                else
                    printf '[%d/%d] Completed: %s\n' "$completed_jobs" "$file_count" "${job_names[0]}" >&2
                fi
                job_pids=("${job_pids[@]:1}")
                job_names=("${job_names[@]:1}")
            fi
        done < <(printf '%s' "$filtered_files" | jq -c '.[]')

        # Drain remaining jobs — every job finishes before we report
        for i in "${!job_pids[@]}"; do
            completed_jobs=$((completed_jobs + 1))
            if ! wait "${job_pids[$i]}"; then
                failed_files+=("${job_names[$i]}")
                printf '[%d/%d] Failed: %s\n' "$completed_jobs" "$file_count" "${job_names[$i]}" >&2
            else
                printf '[%d/%d] Completed: %s\n' "$completed_jobs" "$file_count" "${job_names[$i]}" >&2
            fi
        done

        # Report per-file failures after the batch (same style as the
        # sequential path's "Failed to download" warning)
        for i in "${!failed_files[@]}"; do
            warn "Failed to download ${failed_files[$i]}"
        done

        rm -rf "$temp_dir"
        trap - EXIT
    else
        # Single-file download — the same honesty rules as the parallel path:
        # a failure is recorded in failed_files so the model is never marked
        # complete (.hf-meta suppressed) and the tool exits rc 1.
        while IFS= read -r file_json; do
            local fname fsize
            fname="$(printf '%s' "$file_json" | jq -r '.rfilename')"
            fsize="$(printf '%s' "$file_json" | jq -r '.size // 0')"
            total_size=$((total_size + fsize))

            local url="${HF_BASE}/${ns}/${repo}/resolve/${branch}/${fname}"
            local target="${target_dir}/${fname}"

            if ! hf_download_with_progress "$url" "$target"; then
                failed_files+=("$fname")
                warn "Failed to download $fname"
                continue
            fi
        done < <(printf '%s' "$filtered_files" | jq -c '.[]')
    fi

    # Write metadata — only when every file in the batch succeeded. A partial
    # failure means the model is incomplete; writing .hf-meta would advertise
    # it as complete to `list`/`cache` and hand incomplete weights to
    # `pos ai server start`.
    if [ "${#failed_files[@]}" -eq 0 ]; then
        local meta_file="${target_dir}/.hf-meta"
        local timestamp
        timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
        local file_list
        file_list="$(printf '%s' "$filtered_files" | jq -c '[.[] | .rfilename]')"

        cat > "$meta_file" <<METAEOF
{
  "repo_id": "$repo_id",
  "branch": "$branch",
  "downloaded_at": "$timestamp",
  "files": $file_list
}
METAEOF
    else
        warn "Not writing .hf-meta — ${repo_id} is incomplete (${#failed_files[@]} file(s) failed)"
    fi

    # Summary — honest in both single-file and parallel paths: any failure
    # yields a success/failure count, never a false "downloaded" claim.
    if [ "${#failed_files[@]}" -gt 0 ]; then
        # Honest count: attempted = total files, success = total − failures
        local success_count=$((file_count - ${#failed_files[@]}))
        printf '📥 Downloaded: %s (%d of %d files, %d failed: %s)\n' \
            "$repo_id" "$success_count" "$file_count" "${#failed_files[@]}" "${failed_files[*]}"
        printf '📁 %s/\n' "$target_dir"
    elif [ "$file_count" -eq 1 ]; then
        local fname
        fname="$(printf '%s' "$filtered_files" | jq -r '.[0].rfilename')"
        local fsize
        fsize="$(printf '%s' "$filtered_files" | jq -r '.[0].size // 0')"
        local human_size
        human_size="$(hf_human_size "$fsize")"
        printf '📥 Downloaded: %s/%s (%s)\n' "$repo_id" "$fname" "$human_size"
        printf '📁 %s/%s\n' "$target_dir" "$fname"
    else
        local total_human
        total_human="$(hf_human_size "$total_size")"
        printf '📥 Downloaded: %s (%d files, %s)\n' "$repo_id" "$file_count" "$total_human"
        printf '📁 %s/\n' "$target_dir"
    fi

    # Any failure — single-file or parallel batch — must be detectable by
    # scripts: exit non-zero.
    if [ "${#failed_files[@]}" -gt 0 ]; then
        return 1
    fi
}

cmd_list() {
    [ -d "$HF_DOWNLOAD_DIR" ] || { warn "No models downloaded yet"; return 0; }

    local models=()
    while IFS= read -r dir; do
        [ -d "$dir" ] || continue
        local meta_file="${dir}/.hf-meta"
        [ -f "$meta_file" ] || continue
        models+=("$dir")
    done < <(find "$HF_DOWNLOAD_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort)

    [ ${#models[@]} -gt 0 ] || { warn "No models downloaded yet"; return 0; }

    printf 'Downloaded models (%d):\n' "${#models[@]}"
    for dir in "${models[@]}"; do
        local meta_file="${dir}/.hf-meta"
        local name
        name="$(basename "$dir")"
        local total_size=0
        local date_str
        date_str="$(jq -r '.downloaded_at // "unknown"' "$meta_file" 2>/dev/null | cut -dT -f1)"

        # Calculate total size
        while IFS= read -r file; do
            [ -f "$file" ] || continue
            local size
            size="$(stat -c%s "$file" 2>/dev/null || echo 0)"
            total_size=$((total_size + size))
        done < <(find "$dir" -type f ! -name '.hf-meta' 2>/dev/null)

        local human_size
        human_size="$(hf_human_size "$total_size")"
        printf '  %-50s %s   %s\n' "$name" "$human_size" "$date_str"
    done
}

cmd_remove() {
    local repo_id="${SUBCMD_ARGS[0]:-}"
    [ -n "$repo_id" ] || err "Usage: pos ai hf remove <repo-id>"

    local target_dir
    target_dir="$(hf_repo_dir "$repo_id")"

    [ -d "$target_dir" ] || err "Model not found: $repo_id"

    # Calculate size before removal
    local total_size=0
    while IFS= read -r file; do
        [ -f "$file" ] || continue
        local size
        size="$(stat -c%s "$file" 2>/dev/null || echo 0)"
        total_size=$((total_size + size))
    done < <(find "$target_dir" -type f 2>/dev/null)

    local human_size
    human_size="$(hf_human_size "$total_size")"

    rm -rf "$target_dir"
    printf 'Removed: %s (freed %s)\n' "$repo_id" "$human_size"
}

cmd_info() {
    local repo_id="${SUBCMD_ARGS[0]:-}"
    [ -n "$repo_id" ] || err "Usage: pos ai hf info <repo-id>"

    local ns="${repo_id%%/*}"
    local repo="${repo_id#*/}"

    local info_json
    info_json="$(hf_api "/models/${ns}/${repo}")" || err "Failed to fetch repository info for $repo_id"

    local model_name
    model_name="$(printf '%s' "$info_json" | jq -r '.id')" 
    local downloads
    downloads="$(printf '%s' "$info_json" | jq -r '.downloads // 0')"
    local likes
    likes="$(printf '%s' "$info_json" | jq -r '.likes // 0')"
    local tags
    tags="$(printf '%s' "$info_json" | jq -r '.tags // [] | join(\", \")')"
    local description
    description="$(printf '%s' "$info_json" | jq -r '.description // \"No description\"')"
    local author
    author="$(printf '%s' "$info_json" | jq -r '.author // \"Unknown\"')"
    local created
    created="$(printf '%s' "$info_json" | jq -r '.createdAt // \"Unknown\"')"
    local last_modified
    last_modified="$(printf '%s' "$info_json" | jq -r '.lastModified // \"Unknown\"')"
    local card_data
    card_data="$(printf '%s' "$info_json" | jq -r '.cardData // {}')"
    local pipeline_tag
    pipeline_tag="$(printf '%s' "$info_json" | jq -r '.pipeline_tag // \"Unknown\"')"
    local model_type
    model_type="$(printf '%s' "$info_json" | jq -r '.modelType // \"Unknown\"')"
    local architectures
    architectures="$(printf '%s' "$info_json" | jq -r '.architectures // [] | join(\", \")')"

    printf "Repository: %s\n" "$model_name"
    printf "Author: %s\n" "$author"
    printf "Description: %s\n" "$description"
    printf "Pipeline tag: %s\n" "$pipeline_tag"
    printf "Model type: %s\n" "$model_type"
    printf "Architectures: %s\n" "$architectures"
    printf "Downloads: %s\n" "$downloads"
    printf "Likes: %s\n" "$likes"
    printf "Created: %s\n" "$created"
    printf "Last modified: %s\n" "$last_modified"
    printf "Tags: %s\n" "$tags"
    printf "\n"

    # Show card data if available
    if [ -n "$card_data" ] && [ "$card_data" != "{}" ]; then
        printf "Card data:\n"
        printf '%s' "$card_data" | jq -r 'to_entries[] | "  \(.key): \(.value)"' 2>/dev/null || printf "  (raw data)\n"
    fi
}

cmd_files() {
    local repo_id="${SUBCMD_ARGS[0]:-}"
    [ -n "$repo_id" ] || err "Usage: pos ai hf files <repo-id>"

    local branch
    branch="$(hf_resolve_branch "$repo_id" "$REVISION")"

    local files_json
    files_json="$(hf_repo_files "$repo_id" "$branch")"

    local count
    count="$(printf '%s' "$files_json" | jq 'length')"
    [ "$count" -gt 0 ] || { warn "No files found in $repo_id (branch: $branch)"; return 0; }

    printf 'Files in %s (branch: %s, %d file(s)):\n' "$repo_id" "$branch" "$count"
    printf '%s' "$files_json" | jq -r 'sort_by(.rfilename)[] | [.rfilename, (.size // 0)] | @tsv' | \
        while IFS=$'\t' read -r rpath rsize; do
            printf '  %-60s %s\n' "$rpath" "$(hf_human_size "$rsize")"
        done
}

cmd_cache() {
    local action="${SUBCMD_ARGS[0]:-status}"
    case "$action" in
        status|"") cmd_cache_status ;;
        clear)     cmd_cache_clear ;;
        *)         err "Usage: pos ai hf cache [status|clear]" ;;
    esac
}

# hf_cache_models → one downloaded model dir per line (same discovery as `list`)
hf_cache_models() {
    [ -d "$HF_DOWNLOAD_DIR" ] || return 0
    local dir
    while IFS= read -r dir; do
        [ -d "$dir" ] || continue
        [ -f "$dir/.hf-meta" ] || continue
        printf '%s\n' "$dir"
    done < <(find "$HF_DOWNLOAD_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort)
}

cmd_cache_status() {
    printf 'Cache dir: %s\n' "$HF_DOWNLOAD_DIR"
    local dir model_dirs=() size_total=0 f fsize
    while IFS= read -r dir; do
        [ -n "$dir" ] || continue
        model_dirs+=("$dir")
    done < <(hf_cache_models)
    if [ "${#model_dirs[@]}" -eq 0 ]; then
        printf 'Models:     0 (nothing downloaded yet)\n'
        return 0
    fi
    for dir in "${model_dirs[@]}"; do
        while IFS= read -r f; do
            [ -f "$f" ] || continue
            fsize="$(stat -c%s "$f" 2>/dev/null || echo 0)"
            size_total=$((size_total + fsize))
        done < <(find "$dir" -type f ! -name '.hf-meta' 2>/dev/null)
    done
    printf 'Models:     %d\n' "${#model_dirs[@]}"
    printf 'Size:       %s\n' "$(hf_human_size "$size_total")"
}

cmd_cache_clear() {
    local dir model_dirs=() size_total=0 f fsize yn
    while IFS= read -r dir; do
        [ -n "$dir" ] || continue
        model_dirs+=("$dir")
    done < <(hf_cache_models)
    if [ "${#model_dirs[@]}" -eq 0 ]; then
        printf 'Cache dir: %s\nNo models downloaded yet — nothing to clear\n' "$HF_DOWNLOAD_DIR"
        return 0
    fi

    printf 'The following downloaded models will be removed:\n'
    for dir in "${model_dirs[@]}"; do
        printf '  %s\n' "$(basename "$dir")"
    done

    # Destructive default n, EOF/invalid input denies — same contract as
    # lib/common.sh confirm(). Reads from /dev/tty (like pos-ai-server's model
    # picker) so the tool stays out of the dispatcher's stdin-wrapper tee.
    printf 'Remove all downloaded models? [y/N]: ' >&2
    IFS= read -r yn 2>/dev/null </dev/tty || yn=""
    case "$yn" in
        [Yy]) ;;
        *) echo 'Aborted — nothing removed' >&2; return 0 ;;
    esac

    for dir in "${model_dirs[@]}"; do
        while IFS= read -r f; do
            [ -f "$f" ] || continue
            fsize="$(stat -c%s "$f" 2>/dev/null || echo 0)"
            size_total=$((size_total + fsize))
        done < <(find "$dir" -type f 2>/dev/null)
        rm -rf "$dir"
    done
    printf 'Cache cleared (freed %s)\n' "$(hf_human_size "$size_total")"
}

# ── Dispatch ───────────────────────────────────────────────────
case "$SUBCMD" in
    search)   cmd_search ;;
    download) cmd_download ;;
    list)     cmd_list ;;
    remove)   cmd_remove ;;
    info)     cmd_info ;;
    files)    cmd_files ;;
    cache)    cmd_cache ;;
    *)        err "Unknown subcommand '$SUBCMD' (see --help)" ;;
esac
