#!/usr/bin/env bash
# Usage: bash setup_env.sh [environment-prefix] [all|both|core|benchmark|evo2]
set -euo pipefail

# Keep unrelated per-user Python packages out of dependency resolution/imports.
export PYTHONNOUSERSITE=1
unset PYTHONPATH

project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
environment_name="${1:-rnaseek-release}"
profile="${2:-all}"
if [[ "$profile" != "all" && "$profile" != "both" && "$profile" != "core" && "$profile" != "benchmark" && "$profile" != "evo2" ]]; then
    echo "Profile must be all, both (core + benchmark), core, benchmark, or evo2." >&2
    exit 2
fi
if [[ "$environment_name" == "base" || ! "$environment_name" =~ ^[A-Za-z0-9_-]+$ ]]; then
    echo "Choose a new environment name containing only letters, digits, underscores or hyphens." >&2
    exit 2
fi
if ! command -v conda >/dev/null 2>&1; then
    echo "Install Miniforge or Miniconda, reopen your terminal, then rerun this script." >&2
    exit 2
fi

# A pre-existing environment may contain unrelated work. Never update it here.
conda_base="$(conda info --base)"
environment_list="$(conda env list --json)"
profiles=()
[[ "$profile" == "all" || "$profile" == "both" || "$profile" == "core" ]] && profiles+=(core)
[[ "$profile" == "all" || "$profile" == "both" || "$profile" == "benchmark" ]] && profiles+=(benchmark)
[[ "$profile" == "all" || "$profile" == "evo2" ]] && profiles+=(evo2)
if [[ "$profile" == "all" || "$profile" == "evo2" ]]; then
    "$conda_base/bin/python" -c 'import platform,sys; sys.exit(0 if sys.platform == "linux" and platform.machine() == "x86_64" else "The pinned Evo 2 GPU backend requires Linux x86-64 (WSL2 is an option on Windows). Use profile both for core/encoder environments only.")'
fi
# Check all requested names before creating either environment.
for selected in "${profiles[@]}"; do
    target_name="$environment_name"
    [[ "$selected" == "benchmark" ]] && target_name="$environment_name-benchmark"
    [[ "$selected" == "evo2" ]] && target_name="$environment_name-evo2"
    environment_exists="$("$conda_base/bin/python" -c \
        'import json,sys,pathlib; name=sys.argv[1]; print(int(any(pathlib.Path(p).name == name for p in json.load(sys.stdin)["envs"])))' \
        "$target_name" <<< "$environment_list")"
    if [[ "$environment_exists" == "1" ]]; then
        echo "Requested environment already exists. Supply a new prefix or choose the missing profile." >&2
        exit 2
    fi
done
for selected in "${profiles[@]}"; do
    target_name="$environment_name"
    requirements="$project_dir/requirements.txt"
    constraints="$project_dir/constraints-core.txt"
    if [[ "$selected" == "benchmark" ]]; then
        target_name="$environment_name-benchmark"
        requirements="$project_dir/requirements-benchmark.txt"
        constraints="$project_dir/constraints-benchmark.txt"
    elif [[ "$selected" == "evo2" ]]; then
        target_name="$environment_name-evo2"
        requirements="$project_dir/requirements-evo2.txt"
        constraints="$project_dir/constraints-evo2.txt"
    fi
    conda create --yes --name "$target_name" --override-channels --channel conda-forge python=3.11 pip
    conda env config vars set --name "$target_name" PYTHONNOUSERSITE=1
    conda run --no-capture-output --name "$target_name" python -m pip install -r "$requirements" -c "$constraints"
    if [[ "$selected" == "evo2" ]]; then
        # Official prebuilt wheel matching Python 3.11 / Torch 2.7 / CUDA 12 / CXX11 ABI.
        # Install after Torch; avoid a local CUDA-toolkit source-build requirement.
        conda run --no-capture-output --name "$target_name" python -m pip install -c "$constraints" \
            'https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiTRUE-cp311-cp311-linux_x86_64.whl'
    fi
    conda run --no-capture-output --name "$target_name" python -m pip check
    if [[ "$selected" == "core" ]]; then
        conda run --no-capture-output --name "$target_name" python -c \
            'import torch,transformers,trl,nbformat,nbclient,RNA; print("Core environment ready; GPUs:",torch.cuda.device_count())'
    elif [[ "$selected" == "benchmark" ]]; then
        conda run --no-capture-output --name "$target_name" python -c \
            'import torch,transformers,multimolecule; print("Benchmark environment ready; GPUs:",torch.cuda.device_count())'
    else
        conda run --no-capture-output --name "$target_name" python -c \
            'import torch,evo2; from vortex.model.model import get_block; print("Evo 2 environment ready; GPUs:",torch.cuda.device_count())'
    fi
    printf '\nActivate with: conda activate %s\n' "$target_name"
done