#!/usr/bin/env bash
set -euo pipefail

# download.sh - Download model files and update model.yaml metadata.
#
# This script reads a model.yaml file, downloads the complete model data from
# the specified Hugging Face repository, and then updates the 'files' array
# in the YAML with the paths of the downloaded files.
#
# This approach is more robust than specifying files manually, as it ensures
# the YAML reflects the actual downloaded content.
#
# Usage: ./tools/download.sh models/llama-2-7b/model.yaml

if [ "$#" -ne 1 ]; then
    echo "Usage: $0 <path-to-model.yaml>" >&2
    exit 1
fi

MODEL_YAML="$1"
MODEL_DIR=$(dirname "$MODEL_YAML")

if [ ! -f "$MODEL_YAML" ]; then
    echo "Model YAML not found: $MODEL_YAML" >&2
    exit 1
fi

# Ensure yq is installed
if ! command -v yq &> /dev/null; then
    echo "Error: yq is not installed. Install it with: pip install yq or brew install yq" >&2
    exit 1
fi

# Ensure huggingface-cli is installed
if ! command -v huggingface-cli &> /dev/null; then
    echo "Error: huggingface-cli is not installed. Install it with: pip install huggingface_hub" >&2
    exit 1
fi

echo "Reading metadata from $MODEL_YAML..."

# Create a temporary file to store the updated YAML content
TMP_YAML=$(mktemp)
trap 'rm -f "$TMP_YAML"' EXIT

cp "$MODEL_YAML" "$TMP_YAML"

# Loop over each format and variant to download files
yq -r '.formats[] | . as $format | .variants[] | . as $variant | "\($format.type)\|\($variant.id)\|\($variant.hf_repo)"' "$MODEL_YAML" | while IFS='|' read -r format_type variant_id hf_repo; do
    echo
    echo "Processing variant: $variant_id (format: $format_type) from $hf_repo"

    DEST_PATH="$MODEL_DIR/$variant_id"
    mkdir -p "$DEST_PATH"

    # Check if files are already downloaded by checking for a non-empty directory
    if [ -n "$(ls -A "$DEST_PATH" 2>/dev/null)" ]; then
        echo "[OK] Files for $variant_id already exist in $DEST_PATH. Skipping download."
    else
        repo_id=${hf_repo#https://huggingface.co/}
        echo "[DL] Downloading files for $variant_id from $repo_id..."
        huggingface-cli download "$repo_id" --local-dir "$DEST_PATH" --local-dir-use-symlinks False
    fi

    # After downloading, list the downloaded files relative to the model directory
    downloaded_files=()
    while IFS= read -r file; do
        downloaded_files+=("$(realpath --relative-to="$MODEL_DIR" "$file")")
    done < <(find "$DEST_PATH" -type f)

    # Update the YAML file with the list of downloaded files for the current variant
    echo "Updating $MODEL_YAML with downloaded file paths for $variant_id..."
    # Create a yq expression to update the files for the specific variant
    yq_exp="(.formats[] | select(.type == \"$format_type\") | .variants[] | select(.id == \"$variant_id\") | .files) = []"
    yq eval -i "$yq_exp" "$TMP_YAML"

    for file in "${downloaded_files[@]}"; do
        yq_exp="(.formats[] | select(.type == \"$format_type\") | .variants[] | select(.id == \"$variant_id\") | .files) += [\"$file\"]"
        yq eval -i "$yq_exp" "$TMP_YAML"
    done
done

# Replace the original YAML with the updated one
mv "$TMP_YAML" "$MODEL_YAML"

echo
echo "✅ Download and YAML update complete for $MODEL_YAML."