#!/bin/bash
set -euo pipefail

# Preflight checks for slurm/eval_gguf_eval.sh.
# Uses same env-override pattern as the main script.

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
if [[ -f "${SCRIPT_DIR}/_load_local_env.sh" ]]; then
    # shellcheck source=slurm/_load_local_env.sh
    source "${SCRIPT_DIR}/_load_local_env.sh"
    load_slurm_local_env "$SCRIPT_DIR"
fi

resolve_model_preset() {
    local preset_input preset_key home_dir
    preset_input="$1"
    home_dir="$2"
    preset_key="${preset_input,,}"

    case "$preset_key" in
        glm47flash|glm-4.7-flash|glm47flash-ud-q4_k_xl|glm47flash_ud_q4_k_xl|glm-4.7-flash-ud-q4_k_xl)
            PRESET_CANONICAL="glm47flash-ud-q4_k_xl"
            PRESET_MODEL_DIR="${home_dir}/models/GLM-4.7-Flash-GGUF"
            PRESET_MODEL_FILE="GLM-4.7-Flash-UD-Q4_K_XL.gguf"
            PRESET_MODEL_SOURCE="unsloth/GLM-4.7-Flash-GGUF"
            ;;
        qwen3coder30ba3b-q4_k_m|qwen3-coder-30b-a3b-q4_k_m|qwen3-coder-30b-a3b-instruct-1m-q4_k_m|qwen3coder-q4_k_m|qwen3coder30ba3b)
            PRESET_CANONICAL="qwen3coder30ba3b-q4_k_m"
            PRESET_MODEL_DIR="${home_dir}/models/Qwen3-Coder-30B-A3B-Instruct-1M-GGUF"
            PRESET_MODEL_FILE="Qwen3-Coder-30B-A3B-Instruct-1M-Q4_K_M.gguf"
            PRESET_MODEL_SOURCE="unsloth/Qwen3-Coder-30B-A3B-Instruct-1M-GGUF"
            ;;
        qwen3coder30ba3b-q6_k|qwen3-coder-30b-a3b-q6_k|qwen3-coder-30b-a3b-instruct-1m-q6_k|qwen3coder-q6_k)
            PRESET_CANONICAL="qwen3coder30ba3b-q6_k"
            PRESET_MODEL_DIR="${home_dir}/models/Qwen3-Coder-30B-A3B-Instruct-1M-GGUF"
            PRESET_MODEL_FILE="Qwen3-Coder-30B-A3B-Instruct-1M-Q6_K.gguf"
            PRESET_MODEL_SOURCE="unsloth/Qwen3-Coder-30B-A3B-Instruct-1M-GGUF"
            ;;
        custom|manual)
            PRESET_CANONICAL="custom"
            PRESET_MODEL_DIR="${MODEL_DIR:-${home_dir}/models}"
            PRESET_MODEL_FILE="${MODEL_FILE:-model.gguf}"
            PRESET_MODEL_SOURCE="custom"
            ;;
        *)
            if [[ -n "${MODEL_DIR:-}" && -n "${MODEL_FILE:-}" ]]; then
                PRESET_CANONICAL="custom"
                PRESET_MODEL_DIR="${MODEL_DIR}"
                PRESET_MODEL_FILE="${MODEL_FILE}"
                PRESET_MODEL_SOURCE="custom"
                echo "[WARN] unknown MODEL_PRESET='${preset_input}', falling back to explicit MODEL_DIR/MODEL_FILE."
                return
            fi
            echo "[FAIL] unsupported MODEL_PRESET: ${preset_input}"
            echo "       Supported: glm47flash-ud-q4_k_xl, qwen3coder30ba3b-q4_k_m, qwen3coder30ba3b-q6_k, custom"
            exit 1
            ;;
    esac
}

DEFAULT_HOME_DIR="${HOME:-}"
if [[ -z "$DEFAULT_HOME_DIR" ]]; then
    DEFAULT_HOME_DIR="/home/${USER:-user}"
fi
DEFAULT_REPO_BASENAME="llmdebug"
DEFAULT_MODEL_PRESET="qwen3coder30ba3b-q4_k_m"
DEFAULT_SERVER_CONTAINER="llama-server-cuda"
DEFAULT_EVAL_CONTAINER="python-eval"

HOME_DIR="${HOME_DIR:-$DEFAULT_HOME_DIR}"
REPO_DIR="${REPO_DIR:-${HOME_DIR}/${DEFAULT_REPO_BASENAME}}"
MODEL_PRESET="${MODEL_PRESET:-$DEFAULT_MODEL_PRESET}"
resolve_model_preset "$MODEL_PRESET" "$HOME_DIR"
MODEL_PRESET="$PRESET_CANONICAL"
MODEL_SOURCE="$PRESET_MODEL_SOURCE"
MODEL_DIR="${MODEL_DIR:-$PRESET_MODEL_DIR}"
MODEL_FILE="${MODEL_FILE:-$PRESET_MODEL_FILE}"
SERVER_CONTAINER="${SERVER_CONTAINER:-$DEFAULT_SERVER_CONTAINER}"
EVAL_CONTAINER="${EVAL_CONTAINER:-$DEFAULT_EVAL_CONTAINER}"

FAIL_COUNT=0
WARN_COUNT=0

ok() {
    echo "[OK]   $*"
}

warn() {
    WARN_COUNT=$((WARN_COUNT + 1))
    echo "[WARN] $*"
}

fail() {
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo "[FAIL] $*"
}

command_exists() {
    command -v "$1" >/dev/null 2>&1
}

check_command() {
    local cmd
    cmd="$1"
    if command_exists "$cmd"; then
        ok "command available: $cmd"
    else
        fail "missing command: $cmd"
    fi
}

print_df() {
    local path
    path="$1"
    if [[ -e "$path" ]]; then
        df -h "$path" | sed 's/^/       /'
    fi
}

echo "== GGUF Eval Preflight Check =="
echo "Date:      $(date -Iseconds)"
echo "Host:      $(hostname)"
echo "HOME_DIR:  $HOME_DIR"
echo "REPO_DIR:  $REPO_DIR"
echo "MODEL_PRESET: $MODEL_PRESET"
echo "MODEL_SRC: ${MODEL_SOURCE}"
echo "MODEL_DIR: $MODEL_DIR"
echo "MODEL_FILE:$MODEL_FILE"

echo ""
echo "-- Commands --"
check_command enroot
check_command sbatch
check_command curl
check_command stat
check_command awk

echo ""
echo "-- Paths --"
if REPO_DIR_REAL="$(cd "$REPO_DIR" 2>/dev/null && pwd -P)"; then
    ok "repo directory exists: $REPO_DIR_REAL"
else
    fail "repo directory missing: $REPO_DIR"
fi

if MODEL_DIR_REAL="$(cd "$MODEL_DIR" 2>/dev/null && pwd -P)"; then
    ok "model directory exists: $MODEL_DIR_REAL"
else
    fail "model directory missing: $MODEL_DIR"
fi

if [[ -n "${MODEL_DIR_REAL:-}" ]]; then
    MODEL_PATH="${MODEL_DIR_REAL}/${MODEL_FILE}"
    if [[ -f "$MODEL_PATH" ]]; then
        ok "model file exists: $MODEL_PATH"
    else
        fail "model file missing: $MODEL_PATH"
    fi
fi

if [[ "$MODEL_PRESET" == "custom" ]]; then
    if [[ -z "${MODEL_ALIAS:-}" ]]; then
        warn "MODEL_PRESET=custom should set MODEL_ALIAS for API compatibility."
    else
        ok "custom model alias set: ${MODEL_ALIAS}"
    fi
fi

if [[ -n "${REPO_DIR_REAL:-}" ]]; then
    if [[ -f "${REPO_DIR_REAL}/slurm/eval_gguf_eval.sh" ]]; then
        ok "main script found: ${REPO_DIR_REAL}/slurm/eval_gguf_eval.sh"
    else
        fail "main script missing: ${REPO_DIR_REAL}/slurm/eval_gguf_eval.sh"
    fi

    if [[ -d "${REPO_DIR_REAL}/evals" ]]; then
        ok "evals directory found: ${REPO_DIR_REAL}/evals"
    else
        fail "evals directory missing: ${REPO_DIR_REAL}/evals"
    fi

    for helper in \
        "slurm/run_gguf_smoke.sh" \
        "slurm/submit_gguf_sweep.sh" \
        "slurm/submit_gguf_resume_chain.sh" \
        "slurm/collect_gguf_debug_bundle.sh"; do
        if [[ -x "${REPO_DIR_REAL}/${helper}" ]]; then
            ok "helper script executable: ${REPO_DIR_REAL}/${helper}"
        else
            warn "helper script missing/non-executable: ${REPO_DIR_REAL}/${helper}"
        fi
    done
fi

echo ""
echo "-- Enroot --"
ENROOT_DATA_DIR="${ENROOT_DATA_PATH:-$HOME/.local/share/enroot}"
if ENROOT_DATA_DIR_REAL="$(cd "$ENROOT_DATA_DIR" 2>/dev/null && pwd -P)"; then
    ENROOT_DATA_DIR="$ENROOT_DATA_DIR_REAL"
    ok "enroot data path: $ENROOT_DATA_DIR"
else
    fail "enroot data path missing: $ENROOT_DATA_DIR"
fi

if command_exists enroot; then
    if enroot list | grep -q "^${SERVER_CONTAINER}$"; then
        ok "enroot container exists: ${SERVER_CONTAINER}"
    else
        fail "enroot container missing: ${SERVER_CONTAINER}"
    fi

    if enroot list | grep -q "^${EVAL_CONTAINER}$"; then
        ok "enroot container exists: ${EVAL_CONTAINER}"
    else
        fail "enroot container missing: ${EVAL_CONTAINER}"
    fi
fi

if [[ -n "${ENROOT_DATA_DIR:-}" ]]; then
    if [[ -d "${ENROOT_DATA_DIR}/${SERVER_CONTAINER}" ]]; then
        ok "server rootfs exists: ${ENROOT_DATA_DIR}/${SERVER_CONTAINER}"
    else
        fail "server rootfs missing: ${ENROOT_DATA_DIR}/${SERVER_CONTAINER}"
    fi

    if [[ -d "${ENROOT_DATA_DIR}/${EVAL_CONTAINER}" ]]; then
        ok "eval rootfs exists: ${ENROOT_DATA_DIR}/${EVAL_CONTAINER}"
    else
        fail "eval rootfs missing: ${ENROOT_DATA_DIR}/${EVAL_CONTAINER}"
    fi
fi

echo ""
echo "-- Storage Snapshot --"
if [[ -n "${MODEL_DIR_REAL:-}" ]]; then
    print_df "$MODEL_DIR_REAL"
fi
if [[ -n "${ENROOT_DATA_DIR:-}" ]]; then
    print_df "$ENROOT_DATA_DIR"
fi
print_df "${TMPDIR:-/tmp}"

if [[ -n "${REPO_DIR_REAL:-}" ]]; then
    if [[ -w "$REPO_DIR_REAL" ]]; then
        ok "repo directory is writable: $REPO_DIR_REAL"
    else
        warn "repo directory is not writable: $REPO_DIR_REAL"
    fi
fi

echo ""
echo "Summary: ${FAIL_COUNT} failure(s), ${WARN_COUNT} warning(s)"

if (( FAIL_COUNT > 0 )); then
    echo "Preflight failed. Fix the listed issues before submitting." 
    exit 1
fi

echo "Preflight passed. You can submit the job."
