#!/usr/bin/env bash
# ── ABI Cloud Bootstrap — Stage 2: Full databases ─────────────────────────
# Downloads all auto-fetchable databases for the 7 ABI plugins into
# ${ABI_RESOURCE_ROOT} (default: <repo>/resources/autoplasm).
#
# Tier-1 databases are fetched via `abi setup-resources` (idempotent, atomic).
# Tier-2 manual databases (CARD, COPLA, abricate, blast, plasme, plasx,
# plasmidhostfinder) cannot be auto-fetched reliably, so this script prints
# step-by-step instructions and records them as `manual_required`.
#
# Usage:
#   bash scripts/cloud/02_databases.sh                      # all plugins, tier 1+2
#   bash scripts/cloud/02_databases.sh --plugin metagenomic_plasmid
#   bash scripts/cloud/02_databases.sh --tier 1             # auto only
#   bash scripts/cloud/02_databases.sh --dry-run
#   bash scripts/cloud/02_databases.sh --resource-root /data/abi_dbs
#
# Environment overrides:
#   ABI_RESOURCE_ROOT          DB root (default: <repo>/resources/autoplasm)
#   ABI_MAMBA_ROOT             mamba env root (default: <repo>/.mamba)
#   ABI_RESOURCE_TIMEOUT_SECONDS  per-DB timeout (default: 86400 = 24h)
# ───────────────────────────────────────────────────────────────────────────

source "$(dirname "${BASH_SOURCE[0]}")/libcommon.sh"

DRY_RUN=false
TIER="all"            # 1 | 2 | all
PLUGIN_SUBSET=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --dry-run)       DRY_RUN=true; shift ;;
    --tier)          TIER="$2"; shift 2 ;;
    --plugin)        PLUGIN_SUBSET="$2"; shift 2 ;;
    --resource-root) ABI_RESOURCE_ROOT="$2"; export ABI_RESOURCE_ROOT; shift 2 ;;
    -h|--help)
      sed -n '2,/^$/p' "$0" | sed 's/^# \?//'
      exit 0 ;;
    *) echo "Unknown: $1"; exit 1 ;;
  esac
done

init_log "02_databases"

# ── Disk precheck: full DB set is ~200 GB ─────────────────────────────────
if [[ "${DRY_RUN}" != "true" ]]; then
  require_free_space_gb 250 "${ABI_RESOURCE_ROOT}"
fi

# ── Determine plugins to process ──────────────────────────────────────────
if [[ -n "${PLUGIN_SUBSET}" ]]; then
  IFS=',' read -ra PLUGINS <<< "${PLUGIN_SUBSET}"
else
  PLUGINS=("${ABI_PLUGINS[@]}")
fi

log_step "Stage 2: Download databases for ${#PLUGINS[@]} plugin(s)"
log_info "Resource root : ${ABI_RESOURCE_ROOT}"
log_info "Mamba root    : ${ABI_MAMBA_ROOT}"
log_info "Plugins       : ${PLUGINS[*]}"
log_info "Tier          : ${TIER}"
[[ "${DRY_RUN}" == "true" ]] && log_warn "DRY RUN — no downloads will occur."

mkdir -p "${ABI_RESOURCE_ROOT}"

# ── Generate a minimal cloud config that points resources.root at the
#    user-chosen ABI_RESOURCE_ROOT (the CLI deep-merges with defaults).
CLOUD_CONFIG="${ABI_LOG_DIR}/cloud_config.yaml"
mkdir -p "${ABI_LOG_DIR}"
cat > "${CLOUD_CONFIG}" <<EOF
# Auto-generated by scripts/cloud/02_databases.sh — do not edit.
resources:
  root: ${ABI_RESOURCE_ROOT}
EOF
log_info "Cloud config  : ${CLOUD_CONFIG}"

# ── abi CLI helper ────────────────────────────────────────────────────────
# run_setup_resources <plugin>  — runs abi setup-resources for the plugin.
run_setup_resources() {
  local plugin="$1"
  local abi_bin; abi_bin="$(env_bin autoplasm-base abi 2>/dev/null || echo abi)"
  if [[ "${DRY_RUN}" == "true" ]]; then
    log_info "  [DRY-RUN] ${abi_bin} setup-resources --type ${plugin} --confirm --config ${CLOUD_CONFIG}"
    return 0
  fi
  if ! command -v "${abi_bin}" >/dev/null 2>&1 && [[ ! -x "${abi_bin}" ]]; then
    log_warn "  abi binary not found (${abi_bin}); is 01_envs.sh complete? Skipping."
    return 1
  fi
  log_info "  Running: ${abi_bin} setup-resources --type ${plugin} --confirm --config ${CLOUD_CONFIG}"
  "${abi_bin}" setup-resources --type "${plugin}" --confirm --config "${CLOUD_CONFIG}" 2>&1 \
    || log_warn "  ${plugin}: setup-resources reported failures (see log)."
}

# ── Tier-1: auto-fetchable databases via abi setup-resources ──────────────
tier1() {
  log_step "Tier 1: Auto-fetchable databases"
  for plugin in "${PLUGINS[@]}"; do
    log_info "→ ${plugin}"
    case "${plugin}" in
      metagenomic_plasmid)
        # The flagship pipeline has the most DBs. Run via the CLI which
        # iterates all auto_setup ResourceSpecs atomically.
        run_setup_resources "${plugin}"
        ;;
      amplicon_16s|wgs_bacteria|rnaseq_expression)
        run_setup_resources "${plugin}"
        ;;
      metatranscriptomics)
        # Shares the rnaseq STAR index; only env setup is auto.
        run_setup_resources "${plugin}"
        ;;
      easymetagenome)
        # kraken2/bracken/kneaddata/humann/metaphlan DBs.
        run_setup_resources "${plugin}"
        ;;
      viral_viwrap)
        log_warn "  viral_viwrap: db_dir is manual_required (site-specific). See tier2."
        ;;
    esac
  done
}

# ── Tier-2: manual databases — print instructions, record status ─────────
tier2() {
  log_step "Tier 2: Manual databases — guided instructions"
  local manual_file="${ABI_LOG_DIR}/manual_databases_required.txt"
  : > "${manual_file}"

  manual() {
    local id="$1" url="$2" dest_rel="$3"
    local dest="${ABI_RESOURCE_ROOT}/${dest_rel}"
    echo "-----------------------------------------------" >> "${manual_file}"
    echo "Database: ${id}" >> "${manual_file}"
    echo "URL:      ${url}" >> "${manual_file}"
    echo "Target:   ${dest}" >> "${manual_file}"
    echo "-----------------------------------------------" >> "${manual_file}"
    log_warn "  MANUAL: ${id} → ${url}"
    log_warn "         target: ${dest}"
  }

  # metagenomic_plasmid Tier-2 DBs
  if [[ " ${PLUGINS[*]} " == *" metagenomic_plasmid "* ]]; then
    manual card              "https://card.mcmaster.ca/download"                       "card"
    manual abricate          "https://github.com/tseemann/abricate#databases"          "abricate"
    manual plasme            "https://github.com/ccb-hms/PLASMe/releases"              "plasme/PLASMe_DB"
    manual plasx_annotations "https://github.com/michaelgoldman/PlasX"                 "plasx/data"
    manual plasx_model       "https://github.com/michaelgoldman/PlasX"                 "plasx/data/plasx_model.pkl"
    manual copla_refgraph    "https://zenodo.org/records/10059131"                     "copla/COPLA_DB"
    manual copla_reflist     "https://zenodo.org/records/10059131"                     "copla/COPLA_DB"
    manual blast             "https://ftp.ncbi.nlm.nih.gov/blast/db/"                  "blast"
    manual plasmidhostfinder "https://bitbucket.org/genomicepidemiology/plasmidfinder_db.git" "plasmidhostfinder"
  fi

  # easymetagenome Tier-2
  if [[ " ${PLUGINS[*]} " == *" easymetagenome "* ]]; then
    manual host_db          "KneadData bowtie2 host index (organism-specific)"         "kneaddata/host"
    manual humann_nucleotide "http://huttenhower.sph.harvard.edu/humann"               "humann/chocophlan"
    manual humann_protein   "http://huttenhower.sph.harvard.edu/humann"               "humann/uniref"
  fi

  # viral_viwrap
  if [[ " ${PLUGINS[*]} " == *" viral_viwrap "* ]]; then
    manual viwrap_db        "https://github.com/Finn-Lab/ViWrap"                       "viwrap/db"
    manual viwrap_envs      "ViWrap's own conda envs (see ViWrap docs)"               "viwrap/envs"
  fi

  # rnaseq / metatranscriptomics organism-specific references
  if [[ " ${PLUGINS[*]} " == *" rnaseq_expression "* ]] || [[ " ${PLUGINS[*]} " == *" metatranscriptomics "* ]]; then
    manual genome_index     "STAR index for your organism (build via STAR --runMode genomeGenerate)" "star_index"
    manual annotation_gtf   "GTF annotation for your organism"                        "annotation.gtf"
  fi

  log_info "Manual instructions written to: ${manual_file}"
  log_warn "Review the file and fetch the listed databases by hand, then re-run 03_verify.sh."
}

# ── Run selected tiers ────────────────────────────────────────────────────
case "${TIER}" in
  1) tier1 ;;
  2) tier2 ;;
  all) tier1; tier2 ;;
  *) log_error "Invalid --tier: ${TIER} (use 1, 2, or all)"; exit 1 ;;
esac

# ── Summary ───────────────────────────────────────────────────────────────
log_step "Stage 2 complete"
if [[ "${DRY_RUN}" != "true" ]]; then
  du -sh "${ABI_RESOURCE_ROOT}" 2>/dev/null || true
  mark_sentinel "${DB_SENTINEL}"
  log_info "Sentinel written: ${DB_SENTINEL}"
fi
echo ""
echo "Next: bash scripts/cloud/03_verify.sh"
