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

cd "$(dirname "$0")/../../.."
conda_prefix="${CONDA_INSTALL_DIR:-$HOME/miniforge3}"
conda_bin="${CONDA_EXE:-conda}"
if ! command -v "$conda_bin" >/dev/null 2>&1; then
    # Reuse a previous installation even when the shell has not loaded Conda yet.
    for candidate in "$conda_prefix/bin/conda" "$HOME/miniconda3/bin/conda" "$HOME/anaconda3/bin/conda" /opt/conda/bin/conda; do
        if [[ -x "$candidate" ]]; then
            conda_bin="$candidate"
            break
        fi
    done
fi

if ! command -v "$conda_bin" >/dev/null 2>&1; then
    if [[ "$(uname -s)" != Linux || "$(uname -m)" != x86_64 ]]; then
        echo "This setup requires Linux x86_64 for the FlashAttention wheel." >&2
        exit 1
    fi
    if [[ -e "$conda_prefix" ]]; then
        echo "No usable Conda found in $conda_prefix. Set CONDA_INSTALL_DIR to a new installation path and rerun." >&2
        exit 1
    fi
    echo "Installing Miniforge (Conda) into $conda_prefix ..."
    installer_dir="$(mktemp -d)"
    trap 'rm -rf "$installer_dir"' EXIT
    wget -O "$installer_dir/miniforge.sh" \
        https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh
    bash "$installer_dir/miniforge.sh" -b -p "$conda_prefix"
    conda_bin="$conda_prefix/bin/conda"
fi

# Enable `conda activate nash` in interactive Bash sessions.
"$conda_bin" init bash

conda_envs="$("$conda_bin" env list)"
if awk '$1 == "nash" { found = 1 } END { exit !found }' <<< "$conda_envs"; then
    "$conda_bin" install --yes --name nash python=3.10 pip
else
    "$conda_bin" create --yes --name nash python=3.10 pip
fi

# Always install through the environment's Python, never the system pip.
pip_install() {
    "$conda_bin" run --no-capture-output --name nash python -m pip install "$@"
}

pip_install -c examples/nash_exp/constraints-server-vllm011.txt -e .
pip_install -c examples/nash_exp/constraints-server-vllm011.txt -r examples/nash_exp/requirements-server.txt
pip_install -c examples/nash_exp/constraints-server-vllm011.txt vllm==0.11.0

wget -c https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
pip_install ./flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

printf '\nSetup complete. Run in your terminal:\n  source ~/.bashrc\n  conda activate nash\n'