#!/bin/bash
#SBATCH --job-name=swebench-lite-gen
#SBATCH --partition=normal
#SBATCH --gres=gpu:full:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=12:00:00
#SBATCH --output=swebench-lite-gen-%j.out
#SBATCH --error=swebench-lite-gen-%j.out

set -euo pipefail

# =============================================================================
# SLURM Job: Generate SWE-Bench Lite predictions (baseline + llmdebug)
# using HAICore Enroot + local OpenAI-compatible llama-server endpoint.
# =============================================================================

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

normalize_flag() {
    local var_name raw
    var_name="$1"
    raw="${!var_name}"
    case "${raw,,}" in
        1|true|yes|on)
            printf -v "$var_name" '1'
            ;;
        0|false|no|off)
            printf -v "$var_name" '0'
            ;;
        *)
            echo "ERROR: ${var_name} must be one of: 1/0/true/false/yes/no/on/off"
            echo "       Got: '${raw}'"
            exit 1
            ;;
    esac
}

validate_int_ge() {
    local var_name min value
    var_name="$1"
    min="$2"
    value="${!var_name}"

    if [[ ! "$value" =~ ^[0-9]+$ ]] || (( value < min )); then
        echo "ERROR: ${var_name} must be an integer >= ${min}. Got '${value}'."
        exit 1
    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_ALIAS="GLM-4.7-Flash"
            PRESET_CONTEXT_SIZE=202752
            PRESET_SERVER_PARALLEL=3
            PRESET_OPENAI_TEMPERATURE="1.0"
            PRESET_RUN_PREFIX="swebench-lite-glm47flash"
            ;;
        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_ALIAS="Qwen3-Coder-30B-A3B-Instruct-1M"
            PRESET_CONTEXT_SIZE=196608
            PRESET_SERVER_PARALLEL=3
            PRESET_OPENAI_TEMPERATURE="0.7"
            PRESET_RUN_PREFIX="swebench-lite-qwen3coder-q4km"
            ;;
        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_ALIAS="Qwen3-Coder-30B-A3B-Instruct-1M"
            PRESET_CONTEXT_SIZE=262144
            PRESET_SERVER_PARALLEL=3
            PRESET_OPENAI_TEMPERATURE="0.7"
            PRESET_RUN_PREFIX="swebench-lite-qwen3coder"
            ;;
        custom|manual)
            PRESET_CANONICAL="custom"
            PRESET_MODEL_DIR="${MODEL_DIR:-${home_dir}/models}"
            PRESET_MODEL_FILE="${MODEL_FILE:-model.gguf}"
            PRESET_MODEL_ALIAS="${MODEL_ALIAS:-custom-model}"
            PRESET_CONTEXT_SIZE=8192
            PRESET_SERVER_PARALLEL=2
            PRESET_OPENAI_TEMPERATURE="0.0"
            PRESET_RUN_PREFIX="swebench-lite-custom"
            ;;
        *)
            if [[ -n "${MODEL_DIR:-}" && -n "${MODEL_FILE:-}" ]]; then
                PRESET_CANONICAL="custom"
                PRESET_MODEL_DIR="${MODEL_DIR}"
                PRESET_MODEL_FILE="${MODEL_FILE}"
                PRESET_MODEL_ALIAS="${MODEL_ALIAS:-custom-model}"
                PRESET_CONTEXT_SIZE=8192
                PRESET_SERVER_PARALLEL=2
                PRESET_OPENAI_TEMPERATURE="0.0"
                PRESET_RUN_PREFIX="swebench-lite-custom"
                echo "WARN: Unknown MODEL_PRESET='${preset_input}', using custom model settings."
                return
            fi
            echo "ERROR: Unsupported MODEL_PRESET='${preset_input}'"
            exit 1
            ;;
    esac
}

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

run_generation_once() {
    local -a gen_args
    gen_args=(
        python -m evals.swebench.cli generate-lite
        --run-id "$SBL_RUN_ID"
        --conditions "$SBL_CONDITIONS"
        --dataset-name "$SBL_DATASET_NAME"
        --split "$SBL_SPLIT"
        --max-instances "$SBL_MAX_INSTANCES"
        --openai-base-url "http://localhost:${SERVER_PORT}"
        --openai-model "$MODEL_ALIAS"
        --openai-api-key "local"
        --openai-timeout-sec "$SBL_OPENAI_TIMEOUT_SEC"
        --openai-max-tokens "$SBL_OPENAI_MAX_TOKENS"
        --openai-temperature "$OPENAI_TEMPERATURE"
        --openai-response-mode "$SBL_OPENAI_RESPONSE_MODE"
        --jobs "$SBL_EVAL_JOBS"
    )

    if [[ -n "$SBL_INSTANCE_IDS_FILE" ]]; then
        gen_args+=(--instance-ids-file "$SBL_INSTANCE_IDS_FILE")
    fi
    if [[ -n "$SBL_SNAPSHOT_DIR" ]]; then
        gen_args+=(--snapshot-dir "$SBL_SNAPSHOT_DIR")
    fi
    if (( SBL_RESUME == 1 )); then
        gen_args+=(--resume)
    fi
    if [[ -n "$SBL_SEED" ]]; then
        gen_args+=(--seed "$SBL_SEED")
    fi

    set +e
    enroot start --rw \
        -m "${REPO_DIR_REAL}:/workspace" \
        -e PYTHONPATH=/workspace/src \
        -- "$EVAL_CONTAINER" \
        bash -lc 'cd /workspace && exec "$@"' _ \
        "${gen_args[@]}" 2>&1 | tee "$GEN_LAUNCH_LOG"
    local rc=${PIPESTATUS[0]}
    set -e
    return "$rc"
}

start_llama_server() {
    enroot start --rw \
        -m "${MODEL_DIR_REAL}:/models" \
        -- "$SERVER_CONTAINER" \
        --model "/models/${MODEL_FILE}" \
        --alias "$MODEL_ALIAS" \
        --port "$SERVER_PORT" \
        --host 127.0.0.1 \
        --ctx-size "$CONTEXT_SIZE" \
        --n-gpu-layers "$GPU_LAYERS" \
        --parallel "$SERVER_PARALLEL" \
        --batch-size "$SERVER_BATCH_SIZE" \
        --ubatch-size "$SERVER_UBATCH_SIZE" \
        --temp "$SERVER_TEMPERATURE" \
        --top-p "$SERVER_TOP_P" \
        --min-p "$SERVER_MIN_P" \
        --repeat-penalty "$SERVER_REPEAT_PENALTY" \
        > "$SERVER_LOG" 2>&1 &

    SERVER_PID=$!
}

cleanup() {
    JOB_EXIT_CODE=$?
    if [[ -n "${SERVER_PID:-}" ]]; then
        kill "$SERVER_PID" 2>/dev/null || true
        pkill -f "llama-server.*${SERVER_PORT}" 2>/dev/null || true
        wait "$SERVER_PID" 2>/dev/null || true
    fi
    echo "Finished: $(date -Iseconds) (exit=${JOB_EXIT_CODE})"
    exit "$JOB_EXIT_CODE"
}

# ---- defaults ----
DEFAULT_HOME_DIR="${HOME:-}"
if [[ -z "$DEFAULT_HOME_DIR" ]]; then
    DEFAULT_HOME_DIR="/home/${USER:-user}"
fi
DEFAULT_REPO_BASENAME="llmdebug"
DEFAULT_MODEL_PRESET="glm47flash-ud-q4_k_xl"
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_DIR="${MODEL_DIR:-$PRESET_MODEL_DIR}"
MODEL_FILE="${MODEL_FILE:-$PRESET_MODEL_FILE}"
MODEL_ALIAS="${MODEL_ALIAS:-$PRESET_MODEL_ALIAS}"

SERVER_CONTAINER="${SERVER_CONTAINER:-$DEFAULT_SERVER_CONTAINER}"
EVAL_CONTAINER="${EVAL_CONTAINER:-$DEFAULT_EVAL_CONTAINER}"

CONTEXT_SIZE="${CONTEXT_SIZE:-$PRESET_CONTEXT_SIZE}"
GPU_LAYERS="${GPU_LAYERS:-99}"
SERVER_PARALLEL="${SERVER_PARALLEL:-$PRESET_SERVER_PARALLEL}"
SERVER_BATCH_SIZE="${SERVER_BATCH_SIZE:-2048}"
SERVER_UBATCH_SIZE="${SERVER_UBATCH_SIZE:-512}"
SERVER_TEMPERATURE="${SERVER_TEMPERATURE:-1.0}"
SERVER_TOP_P="${SERVER_TOP_P:-0.95}"
SERVER_MIN_P="${SERVER_MIN_P:-0.01}"
SERVER_REPEAT_PENALTY="${SERVER_REPEAT_PENALTY:-1.0}"
OPENAI_TEMPERATURE="${OPENAI_TEMPERATURE:-$PRESET_OPENAI_TEMPERATURE}"

SBL_DATASET_NAME="${SBL_DATASET_NAME:-princeton-nlp/SWE-bench_Lite}"
SBL_SPLIT="${SBL_SPLIT:-test}"
SBL_CONDITIONS="${SBL_CONDITIONS:-baseline,llmdebug}"
SBL_MAX_INSTANCES="${SBL_MAX_INSTANCES:-0}"
SBL_INSTANCE_IDS_FILE="${SBL_INSTANCE_IDS_FILE:-}"
SBL_SNAPSHOT_DIR="${SBL_SNAPSHOT_DIR:-}"
SBL_OPENAI_TIMEOUT_SEC="${SBL_OPENAI_TIMEOUT_SEC:-180}"
SBL_OPENAI_MAX_TOKENS="${SBL_OPENAI_MAX_TOKENS:-4096}"
SBL_OPENAI_RESPONSE_MODE="${SBL_OPENAI_RESPONSE_MODE:-diff}"
SBL_RESUME="${SBL_RESUME:-1}"
SBL_SEED="${SBL_SEED:-}"
SBL_EVAL_JOBS="${SBL_EVAL_JOBS:-1}"
SBL_SMOKE_MODE="${SBL_SMOKE_MODE:-0}"
SBL_SMOKE_MAX_INSTANCES="${SBL_SMOKE_MAX_INSTANCES:-10}"
SBL_SMOKE_CONTEXT_SIZE="${SBL_SMOKE_CONTEXT_SIZE:-2048}"

MAX_WAIT="${MAX_WAIT:-600}"
POLL_INTERVAL="${POLL_INTERVAL:-5}"

normalize_flag SBL_RESUME
normalize_flag SBL_SMOKE_MODE

validate_int_ge CONTEXT_SIZE 1
validate_int_ge SERVER_PARALLEL 1
validate_int_ge SBL_MAX_INSTANCES 0
validate_int_ge SBL_EVAL_JOBS 1
validate_int_ge SBL_SMOKE_MAX_INSTANCES 1
validate_int_ge MAX_WAIT 1
validate_int_ge POLL_INTERVAL 1

if (( SBL_SMOKE_MODE == 1 )); then
    SBL_MAX_INSTANCES="$SBL_SMOKE_MAX_INSTANCES"
    CONTEXT_SIZE="$SBL_SMOKE_CONTEXT_SIZE"
fi

if [[ -n "${SLURM_JOB_ID:-}" ]]; then
    JOB_ID="$SLURM_JOB_ID"
else
    JOB_ID="manual-$(date +%s)"
fi

if [[ "$JOB_ID" =~ ^[0-9]+$ ]]; then
    SERVER_PORT="$((19000 + (JOB_ID % 1000)))"
else
    SERVER_PORT="19000"
fi

if [[ -z "${SBL_RUN_ID:-}" ]]; then
    if (( SBL_SMOKE_MODE == 1 )); then
        SBL_RUN_ID="smoke-${PRESET_RUN_PREFIX}-${JOB_ID}"
    else
        SBL_RUN_ID="${PRESET_RUN_PREFIX}-${JOB_ID}"
    fi
fi

if ! MODEL_DIR_REAL="$(cd "$MODEL_DIR" 2>/dev/null && pwd -P)"; then
    echo "ERROR: model directory not found: $MODEL_DIR"
    exit 1
fi
if ! REPO_DIR_REAL="$(cd "$REPO_DIR" 2>/dev/null && pwd -P)"; then
    echo "ERROR: repo directory not found: $REPO_DIR"
    exit 1
fi

MODEL_PATH="${MODEL_DIR_REAL}/${MODEL_FILE}"
if [[ ! -f "$MODEL_PATH" ]]; then
    echo "ERROR: model file not found: $MODEL_PATH"
    exit 1
fi

for required_cmd in enroot curl python; do
    if ! command_exists "$required_cmd"; then
        echo "ERROR: required command missing: $required_cmd"
        exit 1
    fi
done

for ctr in "$SERVER_CONTAINER" "$EVAL_CONTAINER"; do
    if ! enroot list | grep -q "^${ctr}$"; then
        echo "ERROR: enroot container not found: $ctr"
        exit 1
    fi
done

SERVER_LOG="${REPO_DIR_REAL}/swebench-llama-server-${JOB_ID}.log"
GEN_LAUNCH_LOG="${REPO_DIR_REAL}/swebench-generate-launch-${JOB_ID}.log"
RUN_CONFIG_LOG="${REPO_DIR_REAL}/swebench-run-config-${JOB_ID}.json"

cat > "$RUN_CONFIG_LOG" <<JSON
{
  "timestamp": "$(date -Iseconds)",
  "job_id": "${JOB_ID}",
  "run_id": "${SBL_RUN_ID}",
  "model_preset": "${MODEL_PRESET}",
  "model_path": "${MODEL_PATH}",
  "model_alias": "${MODEL_ALIAS}",
  "dataset_name": "${SBL_DATASET_NAME}",
  "split": "${SBL_SPLIT}",
  "conditions": "${SBL_CONDITIONS}",
  "max_instances": ${SBL_MAX_INSTANCES},
  "eval_jobs": ${SBL_EVAL_JOBS},
  "resume": ${SBL_RESUME},
  "smoke_mode": ${SBL_SMOKE_MODE},
  "server_port": ${SERVER_PORT},
  "context_size": ${CONTEXT_SIZE}
}
JSON

if [[ -n "$SBL_INSTANCE_IDS_FILE" && ! -f "$SBL_INSTANCE_IDS_FILE" ]]; then
    echo "ERROR: SBL_INSTANCE_IDS_FILE does not exist: $SBL_INSTANCE_IDS_FILE"
    exit 1
fi

if [[ -n "$SBL_SNAPSHOT_DIR" && ! -d "$SBL_SNAPSHOT_DIR" ]]; then
    echo "ERROR: SBL_SNAPSHOT_DIR does not exist: $SBL_SNAPSHOT_DIR"
    exit 1
fi

echo "========================================"
echo "Job ID:          ${JOB_ID}"
echo "Run ID:          ${SBL_RUN_ID}"
echo "Model preset:    ${MODEL_PRESET}"
echo "Model alias:     ${MODEL_ALIAS}"
echo "Model path:      ${MODEL_PATH}"
echo "Dataset:         ${SBL_DATASET_NAME} (${SBL_SPLIT})"
echo "Conditions:      ${SBL_CONDITIONS}"
echo "Max instances:   ${SBL_MAX_INSTANCES}"
echo "Eval jobs:       ${SBL_EVAL_JOBS}"
echo "Server port:     ${SERVER_PORT}"
echo "Context size:    ${CONTEXT_SIZE}"
echo "Smoke mode:      ${SBL_SMOKE_MODE}"
echo "Server log:      ${SERVER_LOG}"
echo "Generate log:    ${GEN_LAUNCH_LOG}"
echo "Run config:      ${RUN_CONFIG_LOG}"
echo "========================================"

echo "Starting llama-server..."
start_llama_server

trap cleanup EXIT

ELAPSED=0
echo "Waiting for llama-server health endpoint..."
while ! curl -sf "http://localhost:${SERVER_PORT}/health" >/dev/null 2>&1; do
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
        echo "ERROR: llama-server exited before readiness."
        tail -n 50 "$SERVER_LOG" || true
        exit 1
    fi
    if (( ELAPSED >= MAX_WAIT )); then
        echo "ERROR: llama-server did not become ready within ${MAX_WAIT}s"
        tail -n 50 "$SERVER_LOG" || true
        exit 1
    fi
    sleep "$POLL_INTERVAL"
    (( ELAPSED += POLL_INTERVAL )) || true
    echo "  waiting... (${ELAPSED}s/${MAX_WAIT}s)"
done

echo "llama-server ready after ${ELAPSED}s"
echo "Running generation..."
if run_generation_once; then
    GEN_EXIT=0
else
    GEN_EXIT=$?
fi

if (( GEN_EXIT != 0 )); then
    echo "ERROR: generation exited with code ${GEN_EXIT}"
    tail -n 80 "$GEN_LAUNCH_LOG" || true
    exit "$GEN_EXIT"
fi

echo "Generation finished successfully."
echo "Run directory: ${REPO_DIR_REAL}/evals/swebench_runs/${SBL_RUN_ID}"
