#!/usr/bin/env bash
set -uo pipefail

# --------------------------------------------------
# Paths and configuration
# --------------------------------------------------

EXP_ROOT="${EXP_ROOT:-/workspace/v223rc_exp2}"
QA_SCRIPT="${QA_SCRIPT:-/workspace/run_all_qa_hotswap.sh}"
LMF_BIN="${LMF_BIN:-lmf}"

LOG_ROOT="${LOG_ROOT:-/workspace/train_logs}"
STATUS_ROOT="${STATUS_ROOT:-/workspace/train_status}"

NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}"
BASE_QA_PORT="${BASE_QA_PORT:-8000}"
# --------------------------------------------------
# Configurations to skip
# --------------------------------------------------

# Skip training for these configs
SKIP_TRAIN=(
    # lr1b{1..5..1}
    # lr2b{1..4..1}
    # lr3b1
    # lr3b3
    # lr3b4
    # lr3b6
    # lr4b{1..3..1}
    # lr5b{1..3..1}
    # lr5b5
    # lr6b{1..3..1}
)

# Skip QA evaluation for these configs
SKIP_EVAL=(
    # lr1b{2..5..1}
    # lr2b{2..4..1}
    # lr3b3
    # lr3b4
    # lr3b6
    # lr4b{2..3..1}
    # lr5b{1..3..1}
    # lr5b5
    # lr6b{1..3..1}
)

contains() {
    local item="$1"
    shift
    for x in "$@"; do
        [[ "$x" == "$item" ]] && return 0
    done
    return 1
}
# This must match model_name_or_path in the training YAML files.
BASE_MODEL_PATH="${BASE_MODEL_PATH:-/workspace/allenai/Olmo-3-1125-32B}"

mkdir -p "$LOG_ROOT" "$STATUS_ROOT"

# --------------------------------------------------
# Validation
# --------------------------------------------------

if [[ ! -d "$EXP_ROOT" ]]; then
    echo "ERROR: Experiment directory does not exist: $EXP_ROOT" >&2
    exit 1
fi

if [[ ! -f "$QA_SCRIPT" ]]; then
    echo "ERROR: QA script not found: $QA_SCRIPT" >&2
    exit 1
fi

if [[ ! -x "$QA_SCRIPT" ]]; then
    echo "ERROR: QA script is not executable: $QA_SCRIPT" >&2
    echo "Run: chmod +x '$QA_SCRIPT'" >&2
    exit 1
fi

if ! command -v "$LMF_BIN" >/dev/null 2>&1; then
    echo "ERROR: Command not found: $LMF_BIN" >&2
    echo "PATH=$PATH" >&2
    exit 1
fi

if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "ERROR: nvidia-smi not found." >&2
    exit 1
fi

if [[ "$NUM_GPUS" -lt 1 ]]; then
    echo "ERROR: No GPUs detected." >&2
    exit 1
fi

echo "Detected GPUs:  $NUM_GPUS"
echo "Experiment root: $EXP_ROOT"
echo "QA script:       $QA_SCRIPT"
echo "Base model:      $BASE_MODEL_PATH"
echo "Logs:            $LOG_ROOT"
echo "Statuses:        $STATUS_ROOT"

# --------------------------------------------------
# Build ordered job list
# --------------------------------------------------

declare -a JOBS=()

for b in {0..7}; do
    for lr in {0..7}; do
        # if [[ "$lr" -eq 3 && "$b" -eq 2 ]]; then
        #     echo "Skipping baseline: lr3b2"
        #     continue
        # fi

        config="lr${lr}b${b}"
        yaml_path="${EXP_ROOT}/${config}.yaml"

        if [[ ! -f "$yaml_path" ]]; then
            echo "WARNING: Missing configuration: $yaml_path" >&2
            # exit 1
            continue
        fi

        JOBS+=("${config}|${yaml_path}")
    done
done

TOTAL_JOBS="${#JOBS[@]}"

echo "Total jobs: $TOTAL_JOBS"

# --------------------------------------------------
# Train and evaluate one configuration
# --------------------------------------------------

run_job() {
    local gpu_id="$1"
    local config="$2"
    local yaml_path="$3"

    local qa_port=$((BASE_QA_PORT + gpu_id))
    local log_file="${LOG_ROOT}/${config}.log"
    local server_log="${LOG_ROOT}/${config}_server.log"

    local running_file="${STATUS_ROOT}/${config}.running"
    local success_file="${STATUS_ROOT}/${config}.success"
    local failed_file="${STATUS_ROOT}/${config}.failed"

    rm -f "$running_file" "$success_file" "$failed_file"

    {
        echo "Config: $config"
        echo "GPU: $gpu_id"
        echo "YAML: $yaml_path"
        echo "Started: $(date --iso-8601=seconds)"
    } >"$running_file"

    echo "[GPU $gpu_id] Starting $config"

    if (
        set -euo pipefail

        # This process and all its children see exactly one GPU.
        export CUDA_VISIBLE_DEVICES="$gpu_id"

        echo "=================================================="
        echo "Config:               $config"
        echo "Physical GPU:         $gpu_id"
        echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
        echo "YAML:                 $yaml_path"
        echo "QA port:              $qa_port"
        echo "Started:              $(date --iso-8601=seconds)"
        echo "=================================================="

        python - <<'PY'
import os
import torch

print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("CUDA available:", torch.cuda.is_available())
print("Visible GPU count:", torch.cuda.device_count())

if torch.cuda.is_available():
    print("Visible GPU name:", torch.cuda.get_device_name(0))
PY

            if contains "$config" "${SKIP_TRAIN[@]}"; then
        echo
        echo "===== Skipping training: $config ====="
    else
        echo
        echo "===== Training: $config ====="
    
        "$LMF_BIN" train "$yaml_path"
    fi
    
    if contains "$config" "${SKIP_EVAL[@]}"; then
        echo
        echo "===== Skipping QA evaluation: $config ====="
    else
        echo
        echo "===== QA evaluation: $config ====="
    
        CONFIGS="$config" \
        BASE_MODEL_PATH="/workspace/allenai/Olmo-3.1-32B-Instruct" \
        PORT="$qa_port" \
        SERVER_URL="http://127.0.0.1:${qa_port}/v1/chat/completions" \
        ADAPTER_ADMIN_URL="http://127.0.0.1:${qa_port}" \
        SERVER_LOG="$server_log" \
        START_SERVER=1 \
            "$QA_SCRIPT"
    fi

        echo
        echo "Completed: $(date --iso-8601=seconds)"
    ) >"$log_file" 2>&1; then
        mv "$running_file" "$success_file"
        echo "[GPU $gpu_id] Completed $config"
        return 0
    else
        local exit_code=$?

        {
            echo "Exit code: $exit_code"
            echo "Failed: $(date --iso-8601=seconds)"
            echo "Main log: $log_file"
            echo "Server log: $server_log"
        } >>"$running_file"

        mv "$running_file" "$failed_file"

        echo "[GPU $gpu_id] FAILED $config" >&2
        echo "[GPU $gpu_id] See: $log_file" >&2
        return "$exit_code"
    fi
}

# --------------------------------------------------
# Persistent worker: one worker per GPU
# --------------------------------------------------

gpu_worker() {
    local gpu_id="$1"
    local job_index
    local job
    local config
    local yaml_path
    local failures=0

    echo "[GPU $gpu_id] Worker started"

    # GPU 0 gets jobs 0, 4, 8, ...
    # GPU 1 gets jobs 1, 5, 9, ...
    # etc.
    for ((job_index = gpu_id;
          job_index < TOTAL_JOBS;
          job_index += NUM_GPUS)); do

        job="${JOBS[$job_index]}"
        IFS='|' read -r config yaml_path <<<"$job"

        if ! run_job "$gpu_id" "$config" "$yaml_path"; then
            failures=$((failures + 1))
        fi
    done

    echo "[GPU $gpu_id] Worker finished; failures=$failures"

    if [[ "$failures" -gt 0 ]]; then
        return 1
    fi
}

# --------------------------------------------------
# Start workers
# --------------------------------------------------

declare -a worker_pids=()

for ((gpu_id = 0; gpu_id < NUM_GPUS; gpu_id++)); do
    gpu_worker "$gpu_id" &
    worker_pids+=("$!")
done

echo
echo "Started ${#worker_pids[@]} GPU workers."
echo "Monitor with: watch -n 1 nvidia-smi"
echo "Follow logs:  tail -F ${LOG_ROOT}/*.log"
echo

overall_status=0

for pid in "${worker_pids[@]}"; do
    if ! wait "$pid"; then
        overall_status=1
    fi
done

# --------------------------------------------------
# Summary
# --------------------------------------------------

success_count="$(
    find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.success' | wc -l
)"

failed_count="$(
    find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.failed' | wc -l
)"

echo
echo "=================================================="
echo "Parallel training and QA finished"
echo "Successful: $success_count"
echo "Failed:     $failed_count"
echo "Total:      $TOTAL_JOBS"
echo "Logs:       $LOG_ROOT"
echo "=================================================="

if [[ "$failed_count" -gt 0 ]]; then
    echo
    echo "Failed configurations:"

    for file in "$STATUS_ROOT"/*.failed; do
        [[ -e "$file" ]] || continue
        basename "$file" .failed
    done
fi

exit "$overall_status"