#!/bin/bash
# ==============================================================================
# cluster_env.sh — Unified cluster environment loader
#
# Load priority:  env vars  >  ~/.lmflow_cluster.conf  >  defaults here
#
# Usage in any script:
#   SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
#   source "${SCRIPT_DIR}/../../scripts/cluster_env.sh"   # adjust relative path
#
# Or with PROJECT_ROOT already set:
#   source "${PROJECT_ROOT}/scripts/cluster_env.sh"
# ==============================================================================

# ---------- 1. Load user config file (if exists) ----------
_LMFLOW_CONF="${LMFLOW_CLUSTER_CONF:-${HOME}/.lmflow_cluster.conf}"
if [ -f "$_LMFLOW_CONF" ]; then
    # shellcheck disable=SC1090
    source "$_LMFLOW_CONF"
fi

# ---------- 2. Helper: auto-detect slurm account/partitions ----------
_lmflow_detect_slurm_account() {
    if ! command -v sacctmgr >/dev/null 2>&1; then
        return 0
    fi
    sacctmgr -nP show assoc where user="$(whoami)" format=account 2>/dev/null \
        | awk -F'|' 'NF {print $1; exit}'
}

_lmflow_detect_gpu_partitions() {
    if ! command -v sinfo >/dev/null 2>&1; then
        return 0
    fi
    # Keep partitions with at least one node and names matching block GPU queues.
    sinfo -h -o "%P %D" 2>/dev/null | awk '
        {
            part = $1
            nodes = $2 + 0
            gsub(/\*/, "", part)
            if (nodes > 0 && part ~ /^(batch_block|backfill_block)[0-9]+$/) {
                if (!(part in seen)) {
                    seen[part] = 1
                    order[++n] = part
                }
            }
        }
        END {
            for (i = 1; i <= n; i++) {
                printf "%s%s", order[i], (i < n ? "," : "")
            }
        }
    '
}

_lmflow_filter_requested_partitions() {
    local requested="$1"
    local available
    local part
    local avail
    local req_parts=()
    local out_parts=()
    local available_parts=()

    available="$(_lmflow_detect_gpu_partitions)"
    if [ -z "$available" ]; then
        echo ""
        return 0
    fi

    IFS=',' read -r -a available_parts <<< "$available"
    IFS=',' read -r -a req_parts <<< "$requested"

    for part in "${req_parts[@]}"; do
        part="${part// /}"
        part="${part//\*/}"
        [ -z "$part" ] && continue
        for avail in "${available_parts[@]}"; do
            if [ "$part" = "$avail" ]; then
                out_parts+=("$part")
                break
            fi
        done
    done

    if [ ${#out_parts[@]} -eq 0 ]; then
        echo ""
        return 0
    fi

    local IFS=','
    echo "${out_parts[*]}"
}

# ---------- 3. Conda helpers ----------
# Cross-install aware: handles envs listed by name ("dvlm  /path/to/envs/dvlm")
# AND envs listed only by full path ("/path/to/envs/dvlm") which happens when the
# env belongs to a different conda installation than the active base.
_lmflow_conda_env_exists() {
    local env_name="$1"
    command -v conda >/dev/null 2>&1 || return 1
    [ -n "$(_lmflow_resolve_conda_path "${env_name}")" ]
}

# Resolve env name → full path from `conda env list`.
# Returns empty string if not found.
_lmflow_resolve_conda_path() {
    local env_name="$1"
    conda env list 2>/dev/null \
        | awk -v target="${env_name}" '
            /^#/ || /^[[:space:]]*$/ { next }
            {
                gsub(/\*/, "")
                # Case 1: named env — col1 is the name, col2 is the path
                if ($1 == target) { for (i=2;i<=NF;i++) if ($i~/^\//) {print $i; exit} }
                # Case 2: path-only listing — col1 is a path ending in /envs/<name>
                for (i=1;i<=NF;i++) {
                    if ($i ~ /\/envs\// && $i ~ "/"target"$") {print $i; exit}
                }
            }'
}

_lmflow_detect_conda_env() {
    # Prefer currently active env if it's not base and still exists.
    if [ -n "${CONDA_DEFAULT_ENV:-}" ] && [ "${CONDA_DEFAULT_ENV}" != "base" ]; then
        if _lmflow_conda_env_exists "${CONDA_DEFAULT_ENV}"; then
            echo "${CONDA_DEFAULT_ENV}"
            return 0
        fi
    fi

    # Default candidates can be overridden by user config/env.
    local candidates="${LMFLOW_CONDA_ENV_CANDIDATES:-dvlm,shiyi-vila}"
    local cands=()
    local cand
    IFS=',' read -r -a cands <<< "${candidates}"
    for cand in "${cands[@]}"; do
        cand="${cand// /}"
        [ -z "${cand}" ] && continue
        if _lmflow_conda_env_exists "${cand}"; then
            echo "${cand}"
            return 0
        fi
    done

    # Final fallback for non-interactive cases where `conda` may not be ready yet.
    echo "dvlm"
}

# ---------- 4. Slurm defaults ----------
# Account: explicit value > auto-detect > fallback default
if [ -z "${LMFLOW_SLURM_ACCOUNT:-}" ] || [ "${LMFLOW_SLURM_ACCOUNT}" = "auto" ]; then
    _detected_account="$(_lmflow_detect_slurm_account)"
    export LMFLOW_SLURM_ACCOUNT="${_detected_account:-nvr_elm_llm}"
else
    export LMFLOW_SLURM_ACCOUNT
fi

# Partition: explicit value (filtered) > auto-detect > fallback default
_requested_partition="${LMFLOW_SLURM_PARTITION:-auto}"
if [ "${_requested_partition}" = "auto" ]; then
    _resolved_partition="$(_lmflow_detect_gpu_partitions)"
else
    _resolved_partition="$(_lmflow_filter_requested_partitions "${_requested_partition}")"
    if [ -z "${_resolved_partition}" ]; then
        _resolved_partition="$(_lmflow_detect_gpu_partitions)"
    fi
fi
export LMFLOW_SLURM_PARTITION="${_resolved_partition:-batch_block1}"
export LMFLOW_SLURM_PARTITION_CPU="${LMFLOW_SLURM_PARTITION_CPU:-cpu_short}"
export LMFLOW_SLURM_GPUS_PER_NODE="${LMFLOW_SLURM_GPUS_PER_NODE:-8}"
export LMFLOW_SLURM_CPUS_PER_TASK="${LMFLOW_SLURM_CPUS_PER_TASK:-16}"
export LMFLOW_SLURM_MEM="${LMFLOW_SLURM_MEM:-1024G}"
export LMFLOW_SLURM_TIME="${LMFLOW_SLURM_TIME:-04:00:00}"

# ---------- 5. Conda defaults ----------
# Conda env is independent from dataset registry choice.
# Priority: explicit value > auto-detect > fallback.
if [ -z "${LMFLOW_CONDA_ENV:-}" ] || [ "${LMFLOW_CONDA_ENV}" = "auto" ]; then
    _resolved_conda_env="$(_lmflow_detect_conda_env)"
    export LMFLOW_CONDA_ENV="${_resolved_conda_env:-dvlm}"
else
    export LMFLOW_CONDA_ENV
fi

# ---------- 6. VILA dataset registry ----------
# VILA_DATASETS and DATASET_BACKEND should be set in ~/.bashrc.
# We only set DATASET_BACKEND default here as a convenience.
export DATASET_BACKEND="${DATASET_BACKEND:-webdataset_multi_modal}"

# ---------- 7. Helper: activate conda ----------
# Call  lmflow_activate_conda  anywhere after sourcing this file.
lmflow_activate_conda() {
    if [ -f "${HOME}/.bashrc" ]; then
        # shellcheck disable=SC1091
        source "${HOME}/.bashrc"
    fi
    if command -v conda >/dev/null 2>&1; then
        if ! _lmflow_conda_env_exists "${LMFLOW_CONDA_ENV}"; then
            _fallback_conda_env="$(_lmflow_detect_conda_env)"
            if [ -n "${_fallback_conda_env}" ] && [ "${_fallback_conda_env}" != "${LMFLOW_CONDA_ENV}" ]; then
                echo "[cluster_env] warning: conda env '${LMFLOW_CONDA_ENV}' not found, fallback to '${_fallback_conda_env}'" >&2
                export LMFLOW_CONDA_ENV="${_fallback_conda_env}"
            fi
        fi
        # Try by name first; if that fails, resolve to full path and retry.
        if ! conda activate "${LMFLOW_CONDA_ENV}" 2>/dev/null; then
            local _full_path
            _full_path="$(_lmflow_resolve_conda_path "${LMFLOW_CONDA_ENV}")"
            if [ "${_full_path}" != "${LMFLOW_CONDA_ENV}" ]; then
                echo "[cluster_env] activating via full path: ${_full_path}" >&2
                conda activate "${_full_path}" || true
            else
                echo "[cluster_env] ERROR: failed to activate conda env '${LMFLOW_CONDA_ENV}'" >&2
            fi
        fi
    fi
}

# ---------- 8. Print summary (opt-in via LMFLOW_VERBOSE=1) ----------
if [[ "${LMFLOW_VERBOSE:-0}" == "1" ]]; then
    echo "[cluster_env] account=${LMFLOW_SLURM_ACCOUNT}  partition=${LMFLOW_SLURM_PARTITION}"
    echo "[cluster_env] conda_env=${LMFLOW_CONDA_ENV}  vila_datasets=${VILA_DATASETS}"
fi
