#!/usr/bin/env bash
# ============================================================================
# setup.sh — Install Firedrake + icepack into a conda environment
#
# This script:
#   1. Creates a conda environment with all native dependencies from conda-forge
#   2. Sets environment variables so Firedrake finds conda's PETSc, MPI, etc.
#   3. pip-installs Firedrake and icepack (skipping binary builds for
#      MPI-linked packages so they compile against conda's libraries)
#
# Usage:
#   bash setup.sh                           # env named "firedrake"
#   bash setup.sh my-env                    # env named "my-env"
#   bash setup.sh ./envs/firedrake          # env at path ./envs/firedrake
#   FIREDRAKE_ENV=~/fd FIREDRAKE_REF=main bash setup.sh
#
# Prerequisites:
#   - conda or mamba installed
#   - internet access (to download packages)
#
# Tested on: Linux x86_64, macOS ARM64
# ============================================================================

set -euo pipefail

# ── Configuration ──────────────────────────────────────────────────────────

# Environment name or path — accepts:
#   - A simple name like "firedrake" → conda env -n firedrake
#   - An absolute/relative path like "./envs/fd" → conda env -p ./envs/fd
# Set via FIREDRAKE_ENV env var, first positional arg, or default "firedrake".
ENV_TARGET="${FIREDRAKE_ENV:-${1:-firedrake}}"

# Detect if it's a path (contains / or starts with .)
if [[ "$ENV_TARGET" == */* ]] || [[ "$ENV_TARGET" == .* ]]; then
    # Resolve to absolute path, creating parent dir if needed
    _parent="$(dirname "$ENV_TARGET")"
    mkdir -p "$_parent" 2>/dev/null || true
    ENV_TARGET="$(cd "$_parent" && pwd)/$(basename "$ENV_TARGET")"
    ENV_IS_PATH=true
    CONDA_ENV_FLAG="-p"
else
    ENV_IS_PATH=false
    CONDA_ENV_FLAG="-n"
fi
ENV_DISPLAY="$ENV_TARGET"
ENV_ACTIVATE="$ENV_TARGET"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Firedrake version: set to a release tag (e.g. "2025.10.2") or "main"
FIREDRAKE_REF="${FIREDRAKE_REF:-${2:-2025.10.2}}"

# Prefer mamba if available (much faster solver)
if command -v mamba &> /dev/null; then
    CONDA=mamba
elif command -v conda &> /dev/null; then
    CONDA=conda
else
    echo "ERROR: Neither conda nor mamba found. Please install one first."
    echo "  See: https://docs.conda.io/en/latest/miniconda.html"
    exit 1
fi

# ── Colors ─────────────────────────────────────────────────────────────────

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
DIM='\033[2m'
NC='\033[0m'  # No color

info()  { echo -e "${BLUE}[INFO]${NC} $*"; }
ok()    { echo -e "${GREEN}[OK]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
fail()  { echo -e "${RED}[FAIL]${NC} $*"; exit 1; }

# ── Live rolling tail ─────────────────────────────────────────────────────
# Runs a command while showing a live, continuously-updating view of the
# last N lines of its combined stdout/stderr. When the command finishes,
# the rolling display is cleared and replaced with a summary.
#
# Usage:  run_live <n_lines> <description> <command...>
# Returns the exit code of the command.

TAIL_LINES=20

run_live() {
    local nlines="$1"; shift
    local desc="$1"; shift

    local logfile
    logfile=$(mktemp /tmp/firedrake-install-XXXXXX.log)

    echo -e "${DIM}─── ${desc} ───${NC}"

    # Run the actual command in the background, directing all output to logfile
    "$@" > "$logfile" 2>&1 &
    local cmd_pid=$!

    # Tail the logfile in real-time, showing last N lines with in-place rewriting
    python3 -u -c "
import sys, os, time

nlines = int(sys.argv[1])
path = sys.argv[2]
pid = int(sys.argv[3])
cols = int(os.environ.get('COLUMNS', 120))

buf = []
shown = 0

def render():
    global shown
    # Move up to overwrite previous output
    if shown > 0:
        sys.stdout.write(f'\033[{shown}A')
    shown = 0
    for b in buf:
        sys.stdout.write(f'\033[2K  {b}\n')
        shown += 1
    sys.stdout.flush()

def pid_alive(p):
    try:
        os.kill(p, 0)
        return True
    except ProcessLookupError:
        return False

with open(path, 'r') as f:
    while True:
        line = f.readline()
        if line:
            line = line.rstrip('\n')[:cols]
            buf.append(line)
            if len(buf) > nlines:
                buf = buf[-nlines:]
            render()
        else:
            # No new data — check if process is still running
            if not pid_alive(pid):
                # Drain any remaining lines
                for line in f:
                    line = line.rstrip('\n')[:cols]
                    buf.append(line)
                    if len(buf) > nlines:
                        buf = buf[-nlines:]
                render()
                break
            time.sleep(0.05)

# Clear the rolling display
if shown > 0:
    sys.stdout.write(f'\033[{shown}A')
    for _ in range(shown):
        sys.stdout.write('\033[2K\n')
    sys.stdout.write(f'\033[{shown}A')
    sys.stdout.flush()
" "$nlines" "$logfile" "$cmd_pid"

    # Collect the exit code
    local exit_code=0
    wait "$cmd_pid" || exit_code=$?

    if [[ $exit_code -ne 0 ]]; then
        warn "${desc} failed (exit code ${exit_code}). Full log: ${logfile}"
        echo -e "${DIM}─── Last 30 lines of ${logfile} ───${NC}"
        tail -30 "$logfile"
        echo -e "${DIM}────────────────────────────────────${NC}"
    else
        rm -f "$logfile"
    fi

    return $exit_code
}

# ── Step 1: Create conda environment ──────────────────────────────────────

info "Creating conda environment '${ENV_DISPLAY}' from environment.yml ..."

# Check if environment already exists
ENV_EXISTS=false
if [[ "$ENV_IS_PATH" == true ]]; then
    [[ -d "${ENV_TARGET}/conda-meta" ]] && ENV_EXISTS=true
else
    $CONDA env list 2>/dev/null | grep -q "^${ENV_TARGET} " && ENV_EXISTS=true
fi

if [[ "$ENV_EXISTS" == true ]]; then
    warn "Environment '${ENV_DISPLAY}' already exists."
    read -p "  Recreate it? [y/N] " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        $CONDA env remove ${CONDA_ENV_FLAG} "${ENV_TARGET}" -y
    else
        info "Reusing existing environment."
    fi
fi

$CONDA env create ${CONDA_ENV_FLAG} "${ENV_TARGET}" -f "${SCRIPT_DIR}/environment.yml" || \
    $CONDA env update ${CONDA_ENV_FLAG} "${ENV_TARGET}" -f "${SCRIPT_DIR}/environment.yml"

ok "Conda environment created."

# ── Step 2: Activate and detect paths ─────────────────────────────────────

# We need to source conda's shell hook to use `conda activate` in a script
eval "$(conda shell.bash hook)"
conda activate "${ENV_ACTIVATE}"

CONDA_PREFIX="${CONDA_PREFIX}"
info "CONDA_PREFIX = ${CONDA_PREFIX}"

# Detect platform
OS="$(uname -s)"
ARCH="$(uname -m)"
info "Platform: ${OS} ${ARCH}"

if [[ "$OS" == "Darwin" ]]; then
    SOEXT="dylib"
else
    SOEXT="so"
fi

# ── Step 3: Verify conda packages ────────────────────────────────────────

info "Verifying critical conda packages ..."

# Check PETSc
if [[ ! -f "${CONDA_PREFIX}/lib/libpetsc.${SOEXT}" ]]; then
    fail "PETSc shared library not found at ${CONDA_PREFIX}/lib/libpetsc.${SOEXT}"
fi
ok "PETSc found."

# Check MPI compiler wrappers
if ! command -v mpicc &> /dev/null; then
    fail "mpicc not found. MPI compiler wrappers are missing."
fi
ok "MPI compilers found: $(which mpicc)"

# Check HDF5
if [[ ! -f "${CONDA_PREFIX}/lib/libhdf5.${SOEXT}" ]]; then
    fail "HDF5 not found."
fi
ok "HDF5 found."

# ── Step 4: Set environment variables ─────────────────────────────────────

info "Setting environment variables for Firedrake ..."

# PETSc: conda installs in "prefix" mode (no PETSC_ARCH)
export PETSC_DIR="${CONDA_PREFIX}"
export PETSC_ARCH=""

# Compilers: use MPI wrappers from conda
export CC=mpicc
export CXX=mpicxx
export FC=mpifort

# HDF5: tell h5py and Firedrake to use MPI-enabled HDF5
export HDF5_MPI="ON"
export HDF5_DIR="${CONDA_PREFIX}"

# Firedrake recommends single-threaded BLAS
export OMP_NUM_THREADS=1

# PyOP2 JIT: conda's mpicc wraps x86_64-conda-linux-gnu-cc which PyOP2
# doesn't recognise → falls back to generic compiler without -shared/-fPIC
export PYOP2_CFLAGS="-fPIC -O2"
export PYOP2_LDFLAGS="-shared -Wl,-rpath,${CONDA_PREFIX}/lib"

# Tell pip not to use cached wheels for MPI-linked packages
# (they may be linked against the wrong MPI/PETSc)
export PIP_NO_CACHE_DIR=1

info "  PETSC_DIR  = ${PETSC_DIR}"
info "  PETSC_ARCH = (empty — prefix install)"
info "  CC         = ${CC}"
info "  HDF5_MPI   = ${HDF5_MPI}"

# ── Step 5: Pip install Firedrake ecosystem ───────────────────────────────

info "Installing Firedrake and dependencies via pip ..."
info "This may take several minutes (compiling C extensions) ..."

# Purge pip cache to avoid stale binaries linked against wrong libs
pip cache purge 2>/dev/null || true

# Install packages that have C extensions and must link against conda's libs.
# We use --no-binary to force compilation against the conda environment.
# The order matters: dependencies first, then Firedrake, then icepack.

info "  [1/5] Installing build dependencies ..."
# These are needed as build backends / build-time deps by libsupermesh
# and firedrake. We install them now so --no-build-isolation works
# (which ensures C extensions link against conda's MPI/PETSc, not
# some isolated copy).
# Install ALL build backends that Firedrake's dependency tree uses.
# --no-build-isolation means pip won't auto-install these, so we need
# them all present before building anything.
run_live $TAIL_LINES "Installing build backends" \
    pip install \
    setuptools wheel \
    scikit-build-core \
    meson-python meson ninja \
    hatchling hatch-vcs hatch-fancy-pypi-readme \
    flit-core \
    cython \
    pybind11 \
    petsctools
ok "Build dependencies installed."

info "  [2/5] Installing mpi4py, petsc4py, h5py (from source against conda libs) ..."
# These should already be installed by conda, but if pip needs to reinstall
# them (e.g. for Firedrake version constraints), force source builds.
# Usually conda's versions are fine and pip will skip them.

info "  [3/5] Installing libsupermesh ..."
run_live $TAIL_LINES "libsupermesh" \
    pip install --no-build-isolation libsupermesh || {
    warn "libsupermesh pip install failed, trying from git ..."
    run_live $TAIL_LINES "libsupermesh (git)" \
        pip install --no-build-isolation \
        "libsupermesh @ git+https://github.com/firedrakeproject/libsupermesh.git"
}
ok "libsupermesh installed."

info "  [4/5] Installing Firedrake (${FIREDRAKE_REF}) ..."
# Firedrake's pyproject.toml lists dependencies like petsc4py, mpi4py, h5py
# as git URLs.  Conda already provides working versions of these, so we
# generate a pip constraints file that pins them to the conda-installed
# versions.  This prevents pip from cloning & rebuilding petsc4py (which
# fails on release tags due to a setuptools dry_run incompatibility).
#
# Clone first so the user can see git progress (pip hides it entirely).
# Clones live in ~/.firedrake-conda/clones/ so they persist across
# environment recreations (conda env remove + re-create).
CLONE_DIR="${HOME}/.firedrake-conda/clones"
mkdir -p "${CLONE_DIR}"

info "    Cloning Firedrake (${FIREDRAKE_REF}) ..."
if [[ -d "${CLONE_DIR}/firedrake" ]]; then
    info "    (updating existing clone)"
    # Unshallow if needed (shallow clones can't switch tags)
    if [[ -f "${CLONE_DIR}/firedrake/.git/shallow" ]]; then
        git -C "${CLONE_DIR}/firedrake" fetch --progress --unshallow --tags origin 2>&1 || \
        git -C "${CLONE_DIR}/firedrake" fetch --progress --tags origin 2>&1
    else
        git -C "${CLONE_DIR}/firedrake" fetch --progress --tags origin 2>&1
    fi
    git -C "${CLONE_DIR}/firedrake" checkout "${FIREDRAKE_REF}" 2>&1
    if [[ "${FIREDRAKE_REF}" == "main" ]]; then
        git -C "${CLONE_DIR}/firedrake" reset --hard origin/main
    fi
else
    # Full clone for tags; shallow clone for main
    if [[ "${FIREDRAKE_REF}" == "main" ]]; then
        git clone --progress --depth 1 --branch main \
            https://github.com/firedrakeproject/firedrake.git \
            "${CLONE_DIR}/firedrake" 2>&1
    else
        git clone --progress --branch "${FIREDRAKE_REF}" \
            https://github.com/firedrakeproject/firedrake.git \
            "${CLONE_DIR}/firedrake" 2>&1
    fi
fi

# ── Generate pip constraints file ──────────────────────────────────────
# Pin every conda-provided package so pip doesn't try to rebuild them
# from git URLs in Firedrake's pyproject.toml.
info "    Generating pip constraints for conda packages ..."
CONSTRAINTS="${CONDA_PREFIX}/constraints.txt"
python3 - <<'PYEOF' > "${CONSTRAINTS}"
import importlib.metadata, sys

# Packages that conda provides and pip must not rebuild
CONDA_OWNS = [
    "petsc4py", "slepc4py", "mpi4py", "h5py", "numpy", "scipy",
    "sympy", "cython", "pybind11", "pkgconfig", "rtree",
    "scikit-build-core", "hatchling", "flit-core", "setuptools",
]

for pkg in CONDA_OWNS:
    try:
        version = importlib.metadata.version(pkg)
        print(f"{pkg}=={version}")
    except importlib.metadata.PackageNotFoundError:
        pass
PYEOF

info "    Constraints file:"
cat "${CONSTRAINTS}" | sed 's/^/      /'

# ── Patch pyproject.toml to remove conda-provided deps ─────────────────
# Firedrake's release tags may list petsc4py, slepc4py, mpi4py, h5py as
# git URL deps (@ git+https://...) or version pins.  Pip can't satisfy
# direct URL refs from conda, and tries to rebuild them — which fails
# (e.g. petsc4py's setuptools dry_run bug).  We strip these deps since
# conda already provides them.
info "    Patching pyproject.toml to skip conda-provided deps ..."
PYPROJECT="${CLONE_DIR}/firedrake/pyproject.toml"
python3 - "${PYPROJECT}" <<'PYEOF'
import re, sys

path = sys.argv[1]
text = open(path).read()

# Packages conda provides — remove any line in dependencies[] that
# starts with these names (handles version pins AND @ git+... URLs)
CONDA_PKGS = ["petsc4py", "slepc4py", "mpi4py", "h5py"]

lines = text.split("\n")
patched = []
removed = []
for line in lines:
    stripped = line.strip().strip('"').strip("'").strip(",").strip()
    skip = False
    for pkg in CONDA_PKGS:
        # Match "petsc4py...", "  petsc4py...", '"petsc4py...', etc.
        if re.match(rf"^{re.escape(pkg)}(\s|>|<|=|@|;|$)", stripped):
            skip = True
            removed.append(stripped)
            break
    if not skip:
        patched.append(line)

if removed:
    open(path, "w").write("\n".join(patched))
    for r in removed:
        print(f"    Removed: {r}")
else:
    print("    (no conda-provided deps found to remove)")
PYEOF

# ── pip install with constraints ───────────────────────────────────────
info "    pip install Firedrake (from local clone) ..."
run_live $TAIL_LINES "Firedrake" \
    pip install --no-build-isolation -c "${CONSTRAINTS}" "${CLONE_DIR}/firedrake"
ok "Firedrake installed."

info "  [5/5] Installing icepack ..."
info "    Cloning icepack ..."
if [[ -d "${CLONE_DIR}/icepack" ]]; then
    info "    (updating existing clone)"
    git -C "${CLONE_DIR}/icepack" pull --progress 2>&1
else
    git clone --progress --depth 1 \
        https://github.com/icepack/icepack.git \
        "${CLONE_DIR}/icepack" 2>&1
fi

info "    pip install icepack (from local clone) ..."
run_live $TAIL_LINES "icepack" \
    pip install --no-build-isolation -c "${CONSTRAINTS}" "${CLONE_DIR}/icepack"
ok "icepack installed."

# ── Step 6: Install activation script ────────────────────────────────────

info "Installing conda activation/deactivation scripts ..."

ACTIVATE_DIR="${CONDA_PREFIX}/etc/conda/activate.d"
DEACTIVATE_DIR="${CONDA_PREFIX}/etc/conda/deactivate.d"
mkdir -p "${ACTIVATE_DIR}" "${DEACTIVATE_DIR}"

cat > "${ACTIVATE_DIR}/firedrake-env-vars.sh" << 'ACTIVATE_EOF'
#!/usr/bin/env bash
# Set environment variables needed by Firedrake's JIT compiler.
# These are automatically set on `conda activate`.

# Save previous values so we can restore on deactivate
export _FIREDRAKE_OLD_CC="${CC:-}"
export _FIREDRAKE_OLD_CXX="${CXX:-}"
export _FIREDRAKE_OLD_FC="${FC:-}"
export _FIREDRAKE_OLD_PETSC_DIR="${PETSC_DIR:-}"
export _FIREDRAKE_OLD_PETSC_ARCH="${PETSC_ARCH:-}"
export _FIREDRAKE_OLD_HDF5_MPI="${HDF5_MPI:-}"
export _FIREDRAKE_OLD_HDF5_DIR="${HDF5_DIR:-}"
export _FIREDRAKE_OLD_OMP_NUM_THREADS="${OMP_NUM_THREADS:-}"
export _FIREDRAKE_OLD_PYOP2_CFLAGS="${PYOP2_CFLAGS:-}"
export _FIREDRAKE_OLD_PYOP2_LDFLAGS="${PYOP2_LDFLAGS:-}"

# PETSc
export PETSC_DIR="${CONDA_PREFIX}"
export PETSC_ARCH=""

# Compilers — use MPI wrappers so JIT-compiled kernels link against
# the correct MPI and can be loaded alongside PETSc.
export CC=mpicc
export CXX=mpicxx
export FC=mpifort

# HDF5
export HDF5_MPI="ON"
export HDF5_DIR="${CONDA_PREFIX}"

# Firedrake recommends single-threaded BLAS to avoid oversubscription
export OMP_NUM_THREADS=1

# PyOP2 JIT flags — conda's mpicc wraps x86_64-conda-linux-gnu-cc
# which PyOP2 doesn't recognise, so it falls back to a generic compiler
# class that omits -shared and -fPIC.  These env vars inject the flags.
export PYOP2_CFLAGS="-fPIC -O2"
export PYOP2_LDFLAGS="-shared -Wl,-rpath,${CONDA_PREFIX}/lib"
ACTIVATE_EOF

cat > "${DEACTIVATE_DIR}/firedrake-env-vars.sh" << 'DEACTIVATE_EOF'
#!/usr/bin/env bash
# Restore environment variables on `conda deactivate`.

export CC="${_FIREDRAKE_OLD_CC}"
export CXX="${_FIREDRAKE_OLD_CXX}"
export FC="${_FIREDRAKE_OLD_FC}"
export PETSC_DIR="${_FIREDRAKE_OLD_PETSC_DIR}"
export PETSC_ARCH="${_FIREDRAKE_OLD_PETSC_ARCH}"
export HDF5_MPI="${_FIREDRAKE_OLD_HDF5_MPI}"
export HDF5_DIR="${_FIREDRAKE_OLD_HDF5_DIR}"
export OMP_NUM_THREADS="${_FIREDRAKE_OLD_OMP_NUM_THREADS}"
export PYOP2_CFLAGS="${_FIREDRAKE_OLD_PYOP2_CFLAGS}"
export PYOP2_LDFLAGS="${_FIREDRAKE_OLD_PYOP2_LDFLAGS}"

unset _FIREDRAKE_OLD_CC
unset _FIREDRAKE_OLD_CXX
unset _FIREDRAKE_OLD_FC
unset _FIREDRAKE_OLD_PETSC_DIR
unset _FIREDRAKE_OLD_PETSC_ARCH
unset _FIREDRAKE_OLD_HDF5_MPI
unset _FIREDRAKE_OLD_HDF5_DIR
unset _FIREDRAKE_OLD_OMP_NUM_THREADS
unset _FIREDRAKE_OLD_PYOP2_CFLAGS
unset _FIREDRAKE_OLD_PYOP2_LDFLAGS

# Clean up empty vars
[[ -z "$CC" ]] && unset CC
[[ -z "$CXX" ]] && unset CXX
[[ -z "$FC" ]] && unset FC
[[ -z "$PETSC_DIR" ]] && unset PETSC_DIR
[[ -z "$PETSC_ARCH" ]] && unset PETSC_ARCH
[[ -z "$HDF5_MPI" ]] && unset HDF5_MPI
[[ -z "$HDF5_DIR" ]] && unset HDF5_DIR
[[ -z "$OMP_NUM_THREADS" ]] && unset OMP_NUM_THREADS
[[ -z "$PYOP2_CFLAGS" ]] && unset PYOP2_CFLAGS
[[ -z "$PYOP2_LDFLAGS" ]] && unset PYOP2_LDFLAGS
DEACTIVATE_EOF

chmod +x "${ACTIVATE_DIR}/firedrake-env-vars.sh"
chmod +x "${DEACTIVATE_DIR}/firedrake-env-vars.sh"

ok "Activation scripts installed."

# ── Step 7: Verify ────────────────────────────────────────────────────────

info "Running verification ..."

python "${SCRIPT_DIR}/verify.py" && ok "All checks passed!" || {
    warn "Some checks failed. See output above."
    warn "You may need to debug linking issues."
    warn "Try: conda activate ${ENV_ACTIVATE} && python verify.py"
}

# ── Done ──────────────────────────────────────────────────────────────────

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
ok "Installation complete!"
echo ""
echo "  To use Firedrake + icepack:"
echo ""
echo "    conda activate ${ENV_ACTIVATE}"
echo "    python -c 'import firedrake; import icepack; print(\"Ready!\")'"
echo ""
echo "  Environment variables (PETSC_DIR, CC, etc.) are set automatically"
echo "  when you activate the environment."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"