setup_jetson.sh

#!/usr/bin/env bash
# One-shot bootstrap for running Tuffy on a Jetson Orin (Tegra/CUDA, JetPack).
# First run: installs apt build deps, creates the venv, and source-builds
# llama-cpp-python with CUDA (slow, ~20-30 min). Subsequent runs verify the
# existing CUDA build and dependency lockfile are still good and just launch -
# no apt/cmake/rebuild work unless something actually changed. Model weights
# are assumed to be in place already; this script never downloads them.
#
# Arguments are forwarded to main.py, so `bash scripts/setup_jetson.sh --voice`
# starts Tuffy straight in voice mode.
set -Eeuo pipefail

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

cd "$PROJECT_ROOT"

MARKER_FILE=".venv/.tuffy_cuda_ready"
# 3.10 is the one interpreter every jetson-ai-lab wheel set supports. Checked
# against pypi.jetson-ai-lab.io/jp6/{cu126,cu128,cu129}/+simple/onnxruntime-gpu/:
# all three publish cp310, cu129 additionally publishes cp312, and none of them
# publish cp311. So 3.10 is what works regardless of which CUDA set this board
# resolves to below, and 3.11 - what the repo's .python-version asks for, which
# is right on a laptop - would leave this machine with no GPU onnxruntime at
# all. Every uv call below passes --python explicitly, which outranks the
# .python-version file. pyproject's requires-python is >=3.10, so 3.10 is a
# valid resolve target for the rest of the dependency set.
PYTHON_VERSION="3.10"

echo "======================================"
echo " Tuffy Jetson Orin Setup"
echo "======================================"
echo

if [[ -f /etc/nv_tegra_release ]]; then
    cat /etc/nv_tegra_release
else
    echo "ERROR: Jetson environment not detected (/etc/nv_tegra_release missing)."
    exit 1
fi

echo

if ! command -v uv >/dev/null 2>&1; then
    echo "Installing uv..."
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
fi

echo "uv version:"
uv --version
echo

# --- Locate the JetPack CUDA toolkit -----------------------------------
# JetPack 6.x ships different CUDA minor versions depending on release
# (12.2 through 12.6+), so don't hardcode one. Prefer the /usr/local/cuda
# symlink (always points at the active toolkit); fall back to the newest
# /usr/local/cuda-* directory that actually has nvcc.
CUDA_HOME=""

if [[ -x /usr/local/cuda/bin/nvcc ]]; then
    CUDA_HOME=/usr/local/cuda
else
    for candidate in $(ls -d /usr/local/cuda-* 2>/dev/null | sort -V -r); do
        if [[ -x "$candidate/bin/nvcc" ]]; then
            CUDA_HOME="$candidate"
            break
        fi
    done
fi

if [[ -z "$CUDA_HOME" ]]; then
    echo "ERROR: CUDA toolkit not found under /usr/local/cuda*."
    echo "       Install it via 'sudo apt install nvidia-cuda-toolkit' or the"
    echo "       JetPack SDK Manager, then re-run this script."
    exit 1
fi

export CUDA_HOME
export PATH="$CUDA_HOME/bin:$PATH"
export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"

echo "CUDA detected at $CUDA_HOME:"
"$CUDA_HOME/bin/nvcc" --version
echo

# --- Pin every uv invocation to the project venv's interpreter ---------
# `uv pip install/uninstall` does NOT reliably infer the target
# interpreter from an activated venv when a newer system Python (e.g.
# 3.11+) is also present - it can resolve wheel tags (cp311) against
# that system interpreter instead of the venv's cp310, which is exactly
# why onnxruntime-gpu (cp310-only wheels on Jetson) failed to resolve.
# Passing --python explicitly removes the ambiguity everywhere.
VENV_PYTHON="$PROJECT_ROOT/.venv/bin/python"

# llama-cpp-python, pywhispercpp, onnxruntime, and piper-tts must NEVER be
# touched by plain `uv sync` - the default PyPI wheels are CPU-only (or,
# for piper-tts, would drag in a CPU-only onnxruntime dependency) and
# would silently clobber the CUDA-enabled builds we install below,
# forcing a full rebuild on every subsequent run. Both the explicit flag
# (this script's own `uv sync` calls) and the env var (uv reads
# UV_NO_INSTALL_PACKAGE itself, so it also protects any bare `uv sync`
# run directly in this same process/subshell) are set - belt and
# suspenders, since the RC-file guard below only applies to *future*
# shells, not this one.
#
# Note that fastembed (pulled in for elastimem's semantic recall) declares a
# plain "onnxruntime" dependency. Skipping the install here is what lets it
# ride on the CUDA-enabled onnxruntime-gpu build instead: both packages
# provide the same importable `onnxruntime` module, so fastembed gets GPU
# execution for free rather than quietly running its embedder on CPU.
SYNC_FLAGS=(--no-install-package llama-cpp-python --no-install-package pywhispercpp --no-install-package onnxruntime --no-install-package piper-tts --inexact --extra voice)
export UV_NO_INSTALL_PACKAGE="llama-cpp-python,pywhispercpp,onnxruntime,piper-tts"

verify_cuda() {
    .venv/bin/python - <<'PY'
import sys
try:
    import llama_cpp
except ImportError as e:
    print(f"ImportError: {e}", file=sys.stderr)
    sys.exit(1)

gpu_support = False
if hasattr(llama_cpp, "llama_supports_gpu_offload") and llama_cpp.llama_supports_gpu_offload():
    gpu_support = True

try:
    info = llama_cpp.llama_print_system_info()
    if info:
        if isinstance(info, bytes):
            info = info.decode()
        if "CUDA" in info.upper() or "GGML_CUDA" in info.upper():
            gpu_support = True
        print(info)
except Exception as e:
    print(f"Warning: failed to print system info: {e}", file=sys.stderr)

if not gpu_support:
    print("CUDA/GPU backend not detected in llama-cpp-python", file=sys.stderr)
    sys.exit(1)
PY
}

verify_whisper_cuda() {
    .venv/bin/python - <<'PY'
import sys
try:
    import pywhispercpp.model as m
    if 'use_gpu' not in m.ContextParams.__annotations__:
        raise RuntimeError("use_gpu key not found in ContextParams annotations")
    print("pywhispercpp CUDA verification succeeded")
except Exception as e:
    print(f"pywhispercpp CUDA verification failed: {e}", file=sys.stderr)
    sys.exit(1)
PY
}

verify_onnx_gpu() {
    .venv/bin/python - <<'PY'
import sys
try:
    import onnxruntime as ort
    providers = ort.get_available_providers()
    if 'CUDAExecutionProvider' not in providers:
        raise RuntimeError(f"CUDAExecutionProvider not found in {providers}")
    print("onnxruntime CUDA verification succeeded")
except Exception as e:
    print(f"onnxruntime CUDA verification failed: {e}", file=sys.stderr)
    sys.exit(1)
PY
}

verify_piper_tts() {
    .venv/bin/python - <<'PY'
import sys
try:
    from piper import PiperVoice  # noqa: F401
    print("piper-tts import verification succeeded")
except Exception as e:
    print(f"piper-tts verification failed: {e}", file=sys.stderr)
    sys.exit(1)
PY
    # piper-tts also ships a CLI entry point; confirm it's on PATH inside the venv.
    if ! .venv/bin/python -m piper --help >/dev/null 2>&1; then
        echo "piper CLI module check failed" >&2
        return 1
    fi
}

# Imports the modules main.py touches before it ever loads a model, so a
# missing/broken pure-Python dependency surfaces here with a real traceback
# instead of at launch. Deliberately does NOT construct a Session (that would
# load the gguf and cost a minute); it only proves the import graph resolves.
verify_imports() {
    .venv/bin/python - <<'PY'
import sys
try:
    import cv2, PIL, numpy, psutil, yaml, requests, ddgs, mcp   # noqa: F401
    import fastembed, elastimem                                 # noqa: F401
    import src.embedder, src.memory, src.models, src.settings   # noqa: F401
    import src.skills.loader, src.tools.mcp_client              # noqa: F401
    import tuffy_core                                           # noqa: F401
    print("python import graph verification succeeded")
except Exception as e:
    print(f"import verification failed: {type(e).__name__}: {e}", file=sys.stderr)
    sys.exit(1)
PY
}

# --- Local model weights ------------------------------------------------
# src/models/weights/*/ is gitignored (multi-GB binaries), so git never
# carries them and this script deliberately does not fetch them either - on
# this board they are already in place. The check below is informational
# only: it never downloads and never blocks the launch. It exists because
# main.py's fallback to FALLBACK_MODEL is NOT wrapped in a try/except, so if
# these files ever do go missing the failure lands as a traceback at startup,
# and a named warning here is much easier to act on than that.
WEIGHT_FILES=(
    "src/models/weights/qwen3vl-2b-instruct-q4km/Qwen3VL-2B-Instruct-Q4_K_M.gguf"
    "src/models/weights/qwen3vl-2b-instruct-q80/mmproj-Qwen3VL-2B-Instruct-Q8_0.gguf"
)

check_weights() {
    local rel missing=()
    for rel in "${WEIGHT_FILES[@]}"; do
        [[ -s "$rel" ]] || missing+=("$rel")
    done

    if [[ ${#missing[@]} -eq 0 ]]; then
        echo "Local model weights present."
        return 0
    fi

    echo "WARNING: FALLBACK_MODEL (qwen3vl-2b-instruct-q4km) is missing weight files:" >&2
    printf '  %s\n' "${missing[@]}" >&2
    echo "         Offline mode and the no-API-key fallback will fail until these exist." >&2
    return 0
}

# --- Embedder cache -----------------------------------------------------
# Elastimem's semantic recall runs BAAI/bge-small-en-v1.5 through fastembed,
# a one-time ~67MB download that src/embedder.py caches under
# ~/.cache/tuffy/fastembed. main.py prewarms it at startup, which is fine on
# a networked machine but is exactly the thing that cannot happen on an
# air-gapped Jetson. Pull it now, while this script still has the network it
# needed for everything above.
prewarm_embedder() {
    .venv/bin/python - <<'PY'
import sys
try:
    import src.embedder as embedder
    # offline=False: this script only ever runs with a network (it just
    # installed wheels), and False is what permits the one-time fetch.
    # prepare() reports its own state back rather than raising.
    print("embedder prewarm:", embedder.prepare(offline=False))
except Exception as e:
    # Non-fatal: semantic recall degrades to keyword search, and main.py
    # retries the download itself on the next online startup.
    print(f"embedder prewarm skipped: {type(e).__name__}: {e}", file=sys.stderr)
PY
}

# Fingerprint of only what should invalidate the cached CUDA build: the
# version constraints of the four natively-built packages (a bump in any of
# them means a real rebuild) and this script itself (CMAKE flags,
# architecture, wheel index, etc). Deliberately NOT the whole lockfile -
# unrelated dependency bumps (elastimem, mcp, fastembed, ...) must not
# trigger a ~20-30 min llama-cpp-python source rebuild for nothing.
current_fingerprint() {
    grep -E '"(llama-cpp-python|pywhispercpp|piper-tts|onnxruntime)' pyproject.toml 2>/dev/null \
        | cat - "$SCRIPT_DIR/setup_jetson.sh" 2>/dev/null \
        | shasum -a 256 | awk '{print $1}'
}

#
# Fast path: environment already set up and untouched since last success.
#
if [[ -d .venv && -f "$MARKER_FILE" ]]; then
    echo "Existing CUDA-enabled Tuffy environment detected."

    if [[ "$(cat "$MARKER_FILE")" == "$(current_fingerprint)" ]] \
        && verify_cuda >/dev/null 2>&1 \
        && verify_whisper_cuda >/dev/null 2>&1 \
        && verify_onnx_gpu >/dev/null 2>&1 \
        && verify_piper_tts >/dev/null 2>&1 \
        && verify_imports >/dev/null 2>&1; then
        echo "CUDA backend verified, dependencies unchanged."

        if [[ -f uv.lock ]]; then
            uv sync --frozen --python "$VENV_PYTHON" "${SYNC_FLAGS[@]}"
        else
            uv sync --python "$VENV_PYTHON" "${SYNC_FLAGS[@]}"
        fi

        check_weights

        echo
        echo "Launching Tuffy..."
        echo

        exec .venv/bin/python main.py "$@"
    else
        echo "Environment is stale or failed verification. Re-validating/rebuilding."
        echo
    fi
fi

echo "======================================"
echo " Validating system requirements"
echo "======================================"
echo

# Only run apt updates and tool installs if packages are missing.
# libsndfile1 backs soundfile, libportaudio2 backs sounddevice, and
# alsa-utils gives src/voice/audio.py its aplay/arecord fallback path.
if ! command -v cmake >/dev/null 2>&1 \
    || ! command -v ninja >/dev/null 2>&1 \
    || ! dpkg -s libportaudio2 &>/dev/null \
    || ! dpkg -s libsndfile1 &>/dev/null \
    || ! command -v aplay >/dev/null 2>&1; then
    sudo apt update
    sudo apt install -y \
        build-essential \
        cmake \
        ninja-build \
        pkg-config \
        git \
        curl \
        ca-certificates \
        python3-dev \
        python3-pip \
        python3-setuptools \
        libportaudio2 \
        libsndfile1 \
        alsa-utils \
        espeak-ng
else
    echo "Build tools and audio libraries already present. Skipping apt install."
fi

if [[ -d .venv ]]; then
    EXISTING_PY_VER=$(.venv/bin/python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>/dev/null || echo "")
    if [[ "$EXISTING_PY_VER" != "$PYTHON_VERSION" ]]; then
        echo "Existing virtual environment has Python $EXISTING_PY_VER, but Python $PYTHON_VERSION is required."
        echo "Re-creating virtual environment..."
        rm -rf .venv
    fi
fi

if [[ ! -d .venv ]]; then
    echo "Creating virtual environment (Python $PYTHON_VERSION)..."
    uv venv --python "$PYTHON_VERSION"
else
    echo "Virtual environment already exists. Skipping creation."
fi

source .venv/bin/activate

# .env is gitignored and holds ANTHROPIC_API_KEY / GROQ_API_KEY / MAYA_API_KEY
# plus the optional ELASTIMEM_TIER override. main.py reads it directly (no
# shell export needed), and DEFAULT_MODEL is now an Anthropic card, so seed
# the file from the template rather than leaving the user to discover it.
if [[ ! -f .env && -f .env.example ]]; then
    cp .env.example .env
    echo
    echo "Created .env from .env.example - add your API keys to it if you want"
    echo "online models. Tuffy still runs fully offline on the local gguf without them."
fi

echo
echo "======================================"
echo " Synchronizing project dependencies"
echo "======================================"

if [[ -f uv.lock ]]; then
    uv sync --frozen --python "$VENV_PYTHON" "${SYNC_FLAGS[@]}"
else
    uv sync --python "$VENV_PYTHON" "${SYNC_FLAGS[@]}"
fi

echo
echo "======================================"
echo " Validating llama-cpp-python CUDA build"
echo "======================================"

NEED_REBUILD=true

if ! .venv/bin/python -c "import llama_cpp" 2>&1; then
    echo "llama-cpp-python is not installed or cannot be imported."
elif ! verify_cuda; then
    echo "llama-cpp-python is installed, but CUDA/GPU check failed."
else
    NEED_REBUILD=false
fi

if [[ "$NEED_REBUILD" == "true" ]]; then
    echo "llama-cpp-python is missing or lacks CUDA support. Building from source..."
    echo "(This step compiles llama.cpp for Jetson Orin's SM 8.7 GPU and takes a while.)"
    echo

    export CMAKE_ARGS="-DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=87"
    export FORCE_CMAKE=1

    uv pip uninstall --python "$VENV_PYTHON" -y llama-cpp-python || true

    uv pip install \
        --python "$VENV_PYTHON" \
        --force-reinstall \
        --no-cache-dir \
        --no-binary llama-cpp-python \
        "llama-cpp-python>=0.3.32"
else
    echo "llama-cpp-python is already compiled with CUDA. Skipping rebuild."
fi

echo
echo "======================================"
echo " Validating pywhispercpp CUDA build"
echo "======================================"

NEED_WHISPER_REBUILD=true

if ! .venv/bin/python -c "import pywhispercpp" 2>&1; then
    echo "pywhispercpp is not installed or cannot be imported."
elif ! verify_whisper_cuda; then
    echo "pywhispercpp is installed, but CUDA/GPU check failed."
else
    NEED_WHISPER_REBUILD=false
fi

if [[ "$NEED_WHISPER_REBUILD" == "true" ]]; then
    echo "pywhispercpp is missing or lacks CUDA support. Building from source..."
    echo "(This compiles whisper.cpp for Jetson Orin's SM 8.7 GPU.)"
    echo

    export CMAKE_ARGS="-DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=87"
    export GGML_CUDA=1
    export WHISPER_CUDA=1

    uv pip uninstall --python "$VENV_PYTHON" -y pywhispercpp || true

    uv pip install \
        --python "$VENV_PYTHON" \
        --force-reinstall \
        --no-cache-dir \
        --no-binary pywhispercpp \
        "pywhispercpp>=1.5.0"
else
    echo "pywhispercpp is already compiled with CUDA. Skipping rebuild."
fi

echo
echo "======================================"
echo " Validating onnxruntime-gpu for Jetson"
echo "======================================"

# Pick the Jetson AI Lab wheel index. That index does NOT publish one
# directory per CUDA point release - it currently carries jp6/cu126,
# jp6/cu128 and jp6/cu129 only, so the old "strip the dot off nvcc's version"
# scheme produced URLs like .../jp6/cu122 that 404 and took the install down
# with them. Ask the index what it actually has, then choose the newest
# directory that is not newer than this machine's CUDA.
JETSON_INDEX_ROOT="https://pypi.jetson-ai-lab.io"
L4T_RELEASE=$(head -n 1 /etc/nv_tegra_release | grep -o -E "R[0-9]+") || true
case "$L4T_RELEASE" in
    R38) JP_VERSION="jp7" ;;
    R36) JP_VERSION="jp6" ;;
    *)   JP_VERSION="jp6" ;;
esac

CUDA_VER=$("$CUDA_HOME/bin/nvcc" --version | grep -i -o -E "release [0-9]+\.[0-9]+" | cut -d' ' -f2) || true
CUDA_TAG_NUM=$(echo "${CUDA_VER:-12.6}" | tr -d '.')

list_cuda_dirs() {
    curl -sL --max-time 20 "$JETSON_INDEX_ROOT/$1" 2>/dev/null \
        | grep -o -E 'cu[0-9]+' | sort -u -V || true
}

AVAILABLE_CU=$(list_cuda_dirs "$JP_VERSION")

# The index publishes a jp<N> tree only once wheels for that JetPack exist -
# at time of writing jp6 is the only one. Mapping L4T R38 to "jp7" above is
# correct in principle but would build a 404 URL today, so if the tree is
# empty (or unreachable), drop back to jp6 rather than failing the install.
if [[ -z "$AVAILABLE_CU" && "$JP_VERSION" != "jp6" ]]; then
    echo "Note: $JETSON_INDEX_ROOT/$JP_VERSION has no wheels published; falling back to jp6."
    JP_VERSION="jp6"
    AVAILABLE_CU=$(list_cuda_dirs "$JP_VERSION")
fi

CUDA_SUFFIX=""
if [[ -n "$AVAILABLE_CU" ]]; then
    # Newest index dir <= our CUDA. Falls through to the oldest available if
    # this machine's CUDA predates everything published (an older JetPack);
    # verify_onnx_gpu below is what catches it if that wheel won't load.
    for cu in $AVAILABLE_CU; do
        if [[ "${cu#cu}" -le "$CUDA_TAG_NUM" ]]; then
            CUDA_SUFFIX="$cu"
        fi
    done
    if [[ -z "$CUDA_SUFFIX" ]]; then
        CUDA_SUFFIX=$(echo "$AVAILABLE_CU" | head -n 1)
        echo "Note: CUDA $CUDA_VER is older than any published wheel set; trying $CUDA_SUFFIX."
    fi
else
    CUDA_SUFFIX="cu126"
    echo "Note: could not reach $JETSON_INDEX_ROOT to list wheel sets; assuming $CUDA_SUFFIX."
fi

JETSON_PIP_INDEX="$JETSON_INDEX_ROOT/$JP_VERSION/$CUDA_SUFFIX"

# fastembed caps onnxruntime at <1.24 on Python 3.10 (it excludes 1.24.0/1.24.1
# outright on every version). The jetson-ai-lab index carries both 1.23.0 and
# 1.24.0, and an unpinned install grabs the newer one - which then conflicts
# with the fastembed that elastimem's semantic recall depends on. Pin it.
ONNX_SPEC="onnxruntime-gpu<1.24"

NEED_ONNX_GPU=true
if ! .venv/bin/python -c "import onnxruntime" 2>&1; then
    echo "onnxruntime is not installed or cannot be imported."
elif ! verify_onnx_gpu; then
    echo "onnxruntime-gpu check failed. Re-installing..."
else
    NEED_ONNX_GPU=false
fi

if [[ "$NEED_ONNX_GPU" == "true" ]]; then
    echo "Installing GPU-enabled onnxruntime for Jetson from $JETSON_PIP_INDEX..."
    echo

    uv pip uninstall --python "$VENV_PYTHON" -y onnxruntime onnxruntime-gpu || true

    uv pip install \
        --python "$VENV_PYTHON" \
        --extra-index-url "$JETSON_PIP_INDEX" \
        --no-cache-dir \
        "$ONNX_SPEC"
else
    echo "onnxruntime-gpu is already installed and verified with CUDA support."
fi

echo
echo "======================================"
echo " Validating piper-tts"
echo "======================================"

# piper-tts is the Python front-end for the Piper TTS engine. It runs its
# voice models through onnxruntime, so on Jetson we want it to ride on top
# of the CUDA-enabled onnxruntime-gpu we just verified above rather than
# pulling in the CPU-only "onnxruntime" package as a transitive dependency.
NEED_PIPER=true
if ! .venv/bin/python -c "import piper" 2>&1; then
    echo "piper-tts is not installed or cannot be imported."
elif ! verify_piper_tts; then
    echo "piper-tts is installed, but failed verification."
else
    NEED_PIPER=false
fi

if [[ "$NEED_PIPER" == "true" ]]; then
    echo "Installing piper-tts (without pulling in a competing CPU-only onnxruntime)..."
    echo

    uv pip uninstall --python "$VENV_PYTHON" -y piper-tts || true

    # piper-tts's only real runtime deps are "onnxruntime" and "pathvalidate"
    # (confirmed against piper-tts 1.5.0's metadata). --no-deps skips the plain
    # "onnxruntime" pin so it doesn't clobber the CUDA-enabled
    # onnxruntime-gpu build verified above; numpy is already pulled in via
    # the [voice] extra, so pathvalidate is the only thing left to add.
    uv pip install \
        --python "$VENV_PYTHON" \
        --no-cache-dir \
        --no-deps \
        "piper-tts>=1.4.2"

    uv pip install \
        --python "$VENV_PYTHON" \
        --no-cache-dir \
        "pathvalidate>=3,<4"
else
    echo "piper-tts is already installed and verified."
fi

echo
echo "======================================"
echo " Checking model weights"
echo "======================================"

check_weights

echo
echo "======================================"
echo " Caching the memory embedder"
echo "======================================"

prewarm_embedder

echo
echo "======================================"
echo " Final validation"
echo "======================================"

verify_cuda
verify_whisper_cuda
verify_onnx_gpu
verify_piper_tts
verify_imports

current_fingerprint > "$MARKER_FILE"

echo
echo "======================================"
echo " Writing launcher helper"
echo "======================================"

# $SHELL is the user's *login* shell (from /etc/passwd), which doesn't
# always match the shell actually running this terminal (e.g. login shell
# is zsh but the terminal launched bash). Write to every rc file that
# could plausibly be sourced - bash and zsh both, whichever exist - so
# `tuffy` works regardless of which one this session turns out to be.
RC_FILES=()
[[ -f "$HOME/.bashrc" ]] && RC_FILES+=("$HOME/.bashrc")
[[ -f "$HOME/.zshrc" ]] && RC_FILES+=("$HOME/.zshrc")

# Neither exists yet (fresh account) - create the one matching the login
# shell so there's at least one place the block lands.
if [[ ${#RC_FILES[@]} -eq 0 ]]; then
    case "$(basename "${SHELL:-}")" in
        zsh)  RC_FILES=("$HOME/.zshrc") ;;
        *)    RC_FILES=("$HOME/.bashrc") ;;
    esac
fi

for RC_FILE in "${RC_FILES[@]}"; do
    # Strip any previously written block (marked by these sentinels) so
    # re-running the script always refreshes tuffy()/uv() instead of
    # silently keeping a stale version forever.
    if [[ -f "$RC_FILE" ]] && grep -q '# >>> tuffy launcher >>>' "$RC_FILE"; then
        sed -i.bak '/# >>> tuffy launcher >>>/,/# <<< tuffy launcher <<</d' "$RC_FILE"
    fi

    cat >> "$RC_FILE" <<EOF

# >>> tuffy launcher >>>
export TUFFY_HOME="$PROJECT_ROOT"

# Runs in a subshell and calls the venv's interpreter by path instead of
# activating it, so the caller's shell keeps its own cwd and environment -
# there is nothing to deactivate after Tuffy exits. Arguments pass straight
# through, so "tuffy --voice" starts in voice mode.
tuffy() {
    ( cd "\$TUFFY_HOME" && "\$TUFFY_HOME/.venv/bin/python" main.py "\$@" )
}

# "uv sync" must never reinstall llama-cpp-python/pywhispercpp/onnxruntime/
# piper-tts PyPI wheels over the CUDA-aware builds. UV_NO_INSTALL_PACKAGE is
# uv's own env var for this (see "uv sync --help") - set globally rather
# than scoped to \$TUFFY_HOME so it also covers a bare "uv sync" run from a
# script, cron job, or any non-interactive shell that sourced this rc file
# without going through the uv() function below. Harmless for any other
# project, since it only ever affects packages actually named that.
export UV_NO_INSTALL_PACKAGE="llama-cpp-python,pywhispercpp,onnxruntime,piper-tts"

uv() {
    if [[ "\$PWD" == "\$TUFFY_HOME"* && "\$1" == "sync" ]]; then
        command uv sync --python "\$TUFFY_HOME/.venv/bin/python" --no-install-package llama-cpp-python --no-install-package pywhispercpp --no-install-package onnxruntime --no-install-package piper-tts --inexact --extra voice "\${@:2}"
    else
        command uv "\$@"
    fi
}
# <<< tuffy launcher <<<
EOF
done

echo
echo "======================================"
echo " Setup completed successfully"
echo "======================================"
echo
echo "Open a new terminal (or 'source ~/.bashrc') and run 'tuffy' from anywhere."
echo "'tuffy --voice' starts directly in voice mode."

echo
echo "Launching Tuffy..."
echo

exec .venv/bin/python main.py "$@"