#!/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
}

normalize_server_backend() {
    local raw
    raw="$1"
    case "${raw,,}" in
        vllm)
            printf 'vllm\n'
            ;;
        llama_cpp|llama-cpp|llama.cpp|llama)
            printf 'llama_cpp\n'
            ;;
        *)
            echo "[FAIL] unsupported SERVER_BACKEND='${raw}' (expected: vllm|llama_cpp)"
            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_BACKEND="vllm"
DEFAULT_SERVER_CONTAINER_VLLM="vllm-server-cuda"
DEFAULT_SERVER_CONTAINER_LLAMA_CPP="llama-server-cuda"
DEFAULT_EVAL_CONTAINER="python-eval"
DEFAULT_HF_MODEL_ID="Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
DEFAULT_HF_MODEL_ALIAS="qwen3coder30ba3b"

HOME_DIR="${HOME_DIR:-$DEFAULT_HOME_DIR}"
REPO_DIR="${REPO_DIR:-${HOME_DIR}/${DEFAULT_REPO_BASENAME}}"
SERVER_BACKEND="${SERVER_BACKEND:-$DEFAULT_SERVER_BACKEND}"
SERVER_BACKEND="$(normalize_server_backend "$SERVER_BACKEND")"

HF_MODEL_ID="${HF_MODEL_ID:-$DEFAULT_HF_MODEL_ID}"
HF_MODEL_ALIAS="${HF_MODEL_ALIAS:-$DEFAULT_HF_MODEL_ALIAS}"
HF_CACHE_DIR="${HF_CACHE_DIR:-}"
VLLM_DOWNLOAD_DIR="${VLLM_DOWNLOAD_DIR:-}"
VLLM_KV_CACHE_DTYPE="${VLLM_KV_CACHE_DTYPE:-auto}"
VLLM_CALCULATE_KV_SCALES="${VLLM_CALCULATE_KV_SCALES:-0}"
VLLM_MAX_NUM_SEQS="${VLLM_MAX_NUM_SEQS:-1}"
VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-}"

if [[ "$SERVER_BACKEND" == "llama_cpp" ]]; then
    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_LLAMA_CPP}"
else
    MODEL_PRESET="${MODEL_PRESET:-hf-qwen3coder30ba3b-fp8}"
    MODEL_SOURCE="$HF_MODEL_ID"
    MODEL_DIR="${MODEL_DIR:-}"
    MODEL_FILE="${MODEL_FILE:-}"
    SERVER_CONTAINER="${SERVER_CONTAINER:-$DEFAULT_SERVER_CONTAINER_VLLM}"
fi

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 "== Eval Preflight Check =="
echo "Date:      $(date -Iseconds)"
echo "Host:      $(hostname)"
echo "HOME_DIR:  $HOME_DIR"
echo "REPO_DIR:  $REPO_DIR"
echo "BACKEND:   $SERVER_BACKEND"
echo "SERVER_CTR:${SERVER_CONTAINER}"
echo "EVAL_CTR:  ${EVAL_CONTAINER}"
if [[ "$SERVER_BACKEND" == "llama_cpp" ]]; then
    echo "MODEL_PRESET: $MODEL_PRESET"
    echo "MODEL_SRC:    ${MODEL_SOURCE}"
    echo "MODEL_DIR:    $MODEL_DIR"
    echo "MODEL_FILE:   $MODEL_FILE"
else
    echo "HF_MODEL_ID:   ${HF_MODEL_ID}"
    echo "HF_MODEL_ALIAS:${HF_MODEL_ALIAS}"
    echo "HF_CACHE_DIR:  ${HF_CACHE_DIR:-<unset>}"
    echo "VLLM_DOWNLOAD_DIR: ${VLLM_DOWNLOAD_DIR:-<unset>}"
    echo "VLLM_KV_CACHE_DTYPE: ${VLLM_KV_CACHE_DTYPE}"
    echo "VLLM_CALCULATE_KV_SCALES: ${VLLM_CALCULATE_KV_SCALES}"
    echo "VLLM_MAX_NUM_SEQS: ${VLLM_MAX_NUM_SEQS}"
    echo "VLLM_MAX_NUM_BATCHED_TOKENS: ${VLLM_MAX_NUM_BATCHED_TOKENS:-<unset>}"
fi

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 [[ "$SERVER_BACKEND" == "llama_cpp" ]]; then
    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
else
    if [[ -z "${HF_MODEL_ID//[[:space:]]/}" ]]; then
        fail "HF_MODEL_ID must not be empty when SERVER_BACKEND=vllm"
    else
        ok "HF model id configured"
    fi
    if [[ -z "${HF_MODEL_ALIAS//[[:space:]]/}" ]]; then
        fail "HF_MODEL_ALIAS must not be empty when SERVER_BACKEND=vllm"
    else
        ok "HF model alias configured"
    fi

    if [[ -z "${VLLM_KV_CACHE_DTYPE//[[:space:]]/}" ]]; then
        fail "VLLM_KV_CACHE_DTYPE must not be empty when SERVER_BACKEND=vllm"
    else
        ok "vLLM KV cache dtype configured"
    fi

    if [[ ! "$VLLM_MAX_NUM_SEQS" =~ ^[0-9]+$ ]] || (( VLLM_MAX_NUM_SEQS < 1 )); then
        fail "VLLM_MAX_NUM_SEQS must be an integer >= 1 (got: ${VLLM_MAX_NUM_SEQS})"
    else
        ok "vLLM max num seqs is valid: ${VLLM_MAX_NUM_SEQS}"
    fi

    if [[ -n "${VLLM_MAX_NUM_BATCHED_TOKENS:-}" ]]; then
        if [[ ! "$VLLM_MAX_NUM_BATCHED_TOKENS" =~ ^[0-9]+$ ]] || (( VLLM_MAX_NUM_BATCHED_TOKENS < 1 )); then
            fail "VLLM_MAX_NUM_BATCHED_TOKENS must be empty or integer >= 1 (got: ${VLLM_MAX_NUM_BATCHED_TOKENS})"
        else
            ok "vLLM max batched tokens is valid: ${VLLM_MAX_NUM_BATCHED_TOKENS}"
        fi
    fi

    if [[ -n "$HF_CACHE_DIR" ]]; then
        if HF_CACHE_DIR_REAL="$(cd "$HF_CACHE_DIR" 2>/dev/null && pwd -P)"; then
            ok "HF cache dir exists: $HF_CACHE_DIR_REAL"
            if [[ ! -w "$HF_CACHE_DIR_REAL" ]]; then
                warn "HF cache dir is not writable: $HF_CACHE_DIR_REAL"
            fi
        else
            warn "HF cache dir does not exist yet: $HF_CACHE_DIR"
        fi
    fi

    if [[ -n "$VLLM_DOWNLOAD_DIR" ]]; then
        if VLLM_DOWNLOAD_DIR_REAL="$(cd "$VLLM_DOWNLOAD_DIR" 2>/dev/null && pwd -P)"; then
            ok "vLLM download dir exists: $VLLM_DOWNLOAD_DIR_REAL"
            if [[ ! -w "$VLLM_DOWNLOAD_DIR_REAL" ]]; then
                warn "vLLM download dir is not writable: $VLLM_DOWNLOAD_DIR_REAL"
            fi
        else
            warn "vLLM download dir does not exist yet: $VLLM_DOWNLOAD_DIR"
        fi
    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 "${HF_CACHE_DIR_REAL:-}" ]]; then
    print_df "$HF_CACHE_DIR_REAL"
fi
if [[ -n "${VLLM_DOWNLOAD_DIR_REAL:-}" ]]; then
    print_df "$VLLM_DOWNLOAD_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."
