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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# shellcheck disable=SC1091
source "${SCRIPT_DIR}/catalog-lib.sh"

if [[ $# -lt 1 ]]; then
  echo "Usage: $0 <catalog-model-id>"
  echo ""
  echo "Available models:"
  catalog_list_ids | sed 's/^/  /'
  exit 1
fi

MODEL_ID="$1"

if ! catalog_list_ids | grep -qx "${MODEL_ID}"; then
  echo "ERROR: Unknown model id: ${MODEL_ID}"
  echo "Run: ${SCRIPT_DIR}/list-models.sh"
  exit 1
fi

load_comfyui_stack_env

DATA_ROOT="${DATA_ROOT:-/data}"
catalog_ensure_dirs "${DATA_ROOT}"

SOURCE="$(catalog_get_field "${MODEL_ID}" source)"
NAME="$(catalog_get_field "${MODEL_ID}" name)"
DEST="$(catalog_dest_path "${MODEL_ID}" "${DATA_ROOT}")"

echo "=== Download: ${NAME} (${MODEL_ID}) ==="
echo "Source: ${SOURCE}"
echo "Target: ${DEST}"
echo ""

if [[ -f "${DEST}" ]]; then
  echo "Already exists: ${DEST}"
  echo ""
  echo "Done."
  exit 0
fi

mkdir -p "$(dirname "${DEST}")"

# shellcheck disable=SC1091
source "${SCRIPT_DIR}/download-http-lib.sh"

case "${SOURCE}" in
  huggingface)
    HF_REPO="$(catalog_get_field "${MODEL_ID}" hf_repo)"
    FILENAME="$(catalog_get_field "${MODEL_ID}" filename)"
    DOWNLOAD_URL="$(catalog_get_field "${MODEL_ID}" download_url)"
    HF_REVISION="$(catalog_get_field "${MODEL_ID}" hf_revision)"
    REVISION="${HF_REVISION:-main}"

    if [[ -n "${DOWNLOAD_URL}" ]]; then
      URL="${DOWNLOAD_URL}"
    elif [[ -n "${HF_REPO}" && -n "${FILENAME}" ]]; then
      URL="https://huggingface.co/${HF_REPO}/resolve/${REVISION}/${FILENAME}"
    else
      echo "ERROR: Catalog entry missing download_url or hf_repo/filename"
      exit 1
    fi

    download_http_file "${URL}" "${DEST}"
    ;;
  url)
    DOWNLOAD_URL="$(catalog_get_field "${MODEL_ID}" download_url)"
    if [[ -z "${DOWNLOAD_URL}" ]]; then
      echo "ERROR: Catalog entry missing download_url"
      exit 1
    fi
    download_http_file "${DOWNLOAD_URL}" "${DEST}"
    ;;
  civitai)
    CIVITAI_MODEL_ID="$(catalog_get_field "${MODEL_ID}" civitai_model_id)"
    if [[ -z "${CIVITAI_MODEL_ID}" ]]; then
      echo "ERROR: Catalog entry missing civitai_model_id"
      exit 1
    fi
    if [[ -z "${CIVITAI_TOKEN:-}" ]]; then
      echo "ERROR: CIVITAI_TOKEN not set in stacks/comfyui/.env"
      exit 1
    fi
    curl -fL \
      -H "Authorization: Bearer ${CIVITAI_TOKEN}" \
      -o "${DEST}" \
      "https://civitai.com/api/download/models/${CIVITAI_MODEL_ID}"
    ;;
  *)
    echo "ERROR: Unknown source: ${SOURCE}"
    exit 1
    ;;
esac

if [[ ! -f "${DEST}" ]]; then
  echo "ERROR: Download finished but file missing: ${DEST}"
  exit 1
fi

echo ""
echo "Done."