#!/usr/bin/env bash
# Shared helpers for comfyui-models.catalog.yaml (simple parser, no PyYAML required).

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SERVER_UI_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
CATALOG_FILE="${CATALOG_FILE:-${SERVER_UI_DIR}/comfyui-models.catalog.yaml}"

catalog_list_ids() {
  grep '^  - id:' "${CATALOG_FILE}" | awk '{print $3}'
}

catalog_get_field() {
  local model_id="$1"
  local field="$2"
  awk -v id="${model_id}" -v key="${field}:" '
    $0 ~ "^  - id: " id "$" { found=1; next }
    found && $0 ~ "^  - id:" { exit }
    found && index($0, key) == 5 { sub(/^    [^:]+: /, ""); print; exit }
  ' "${CATALOG_FILE}"
}

catalog_category_dir() {
  local category="$1"
  case "${category}" in
    checkpoint) echo "checkpoints" ;;
    lora) echo "loras" ;;
    vae) echo "vae" ;;
    controlnet) echo "controlnet" ;;
    upscale) echo "upscale_models" ;;
    latent_upscale) echo "latent_upscale_models" ;;
    text_encoder) echo "text_encoders" ;;
    diffusion_models) echo "diffusion_models" ;;
    clip) echo "clip" ;;
    unet) echo "unet" ;;
    *) echo "${category}" ;;
  esac
}

catalog_dest_path() {
  local model_id="$1"
  local data_root="${2:-${DATA_ROOT:-/data}}"
  local local_path category filename subdir
  local_path="$(catalog_get_field "${model_id}" local_path)"
  if [[ -n "${local_path}" ]]; then
    echo "${local_path}"
    return
  fi
  category="$(catalog_get_field "${model_id}" category)"
  filename="$(catalog_get_field "${model_id}" filename)"
  subdir="$(catalog_category_dir "${category}")"
  echo "${data_root}/apps/comfyui/models/${subdir}/${filename}"
}

catalog_ensure_dirs() {
  local data_root="${1:-/data}"
  mkdir -p \
    "${data_root}/apps/comfyui/storage" \
    "${data_root}/apps/comfyui/models/checkpoints" \
    "${data_root}/apps/comfyui/models/loras" \
    "${data_root}/apps/comfyui/models/vae" \
    "${data_root}/apps/comfyui/models/controlnet" \
    "${data_root}/apps/comfyui/models/upscale_models" \
    "${data_root}/apps/comfyui/models/latent_upscale_models" \
    "${data_root}/apps/comfyui/models/text_encoders" \
    "${data_root}/apps/comfyui/models/diffusion_models" \
    "${data_root}/apps/comfyui/models/clip" \
    "${data_root}/apps/comfyui/models/unet" \
    "${data_root}/apps/comfyui/cache/hf-hub" \
    "${data_root}/apps/comfyui/cache/torch-hub" \
    "${data_root}/apps/comfyui/input" \
    "${data_root}/apps/comfyui/output" \
    "${data_root}/apps/comfyui/custom_nodes" \
    "${data_root}/apps/comfyui/workflows"
}

catalog_model_downloaded() {
  local model_id="$1"
  local dest
  dest="$(catalog_dest_path "${model_id}")"
  [[ -f "${dest}" ]]
}

load_comfyui_stack_env() {
  local repo_root
  repo_root="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
  local env_file="${repo_root}/stacks/comfyui/.env"
  if [[ -f "${env_file}" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "${env_file}"
    set +a
  fi
}