ARCH=`uname -m`
if [[ $ARCH == "arm64" ]]; then
    PREFIX="/opt/homebrew"
else
    PREFIX="/usr/local"
fi

export WORKON_HOME=~/python-envs
export VIRTUALENVWRAPPER_VIRTUALENV=$PREFIX/bin/virtualenv
mkdir -p $WORKON_HOME
source $PREFIX/bin/virtualenvwrapper.sh

# NOTE: an inherited VIRTUAL_ENV is *not* re-activated here — that happens at the
# very end of .zshrc so activate captures the fully-built PATH (see the comment
# there). Sourcing virtualenvwrapper.sh above only defines workon/deactivate/etc.

# Mirror the current venv into a per-pane tmux option so that splitting a pane
# can carry it into the new pane (.tmux.conf reads #{@venv} with `-e`). Per-pane
# (not session) state is what lets two panes hold different venvs and lets new
# *windows* start clean. Only fork tmux when the value actually changes.
_venv_track_precmd() {
    [[ -n $TMUX ]] || return
    if [[ "${VIRTUAL_ENV:-}" != "${_LAST_TRACKED_VENV-unset}" ]]; then
        tmux set -p @venv "${VIRTUAL_ENV:-}" 2>/dev/null
        _LAST_TRACKED_VENV="${VIRTUAL_ENV:-}"
    fi
}
autoload -Uz add-zsh-hook
add-zsh-hook precmd _venv_track_precmd

# Which venv an agent shell should use: nearest .venv/venv walking up from $1
# (default $PWD, stopping above $HOME), else the shared fallback env.
_agent_venv_path() {
    local d=${1:-$PWD} c
    while [[ $d == */* && $d != $HOME ]]; do
        for c in "$d/.venv" "$d/venv"; do
            [[ -f $c/bin/activate ]] && { print -r -- "$c"; return 0 }
        done
        d=${d%/*}
    done
    [[ -f $HOME/python-envs/claude/bin/activate ]] && print -r -- "$HOME/python-envs/claude"
}

# Agent shells source this file once per session, rooted at the project dir, and
# snapshot the result: CLAUDECODE marks Claude Code, CODEX_SHELL the ChatGPT/Codex
# desktop app. Only VIRTUAL_ENV is set here; the block at the end of .zshrc turns it
# into an activation once PATH is fully built. Terminal launches use the wrappers below.
if [[ -n ${CLAUDECODE:-} || -n ${CODEX_SHELL:-} ]] && [[ -z ${VIRTUAL_ENV:-} ]]; then
    _venv=$(_agent_venv_path)
    [[ -n $_venv ]] && export VIRTUAL_ENV=$_venv
    unset _venv
fi

# Claude Code takes PATH from its own process env: its shell snapshot writes
# `export PATH=` last, discarding whatever .zshrc built. So the venv has to be
# active *before* launch. The subshell keeps this shell's PATH untouched.
claude() {
    if [[ -n ${VIRTUAL_ENV:-} ]]; then
        command claude "$@"
        return
    fi
    local v=$(_agent_venv_path)
    ( [[ -n $v ]] && source "$v/bin/activate"; exec command claude "$@" )
}

# Codex captures the profile's own PATH, so hand it VIRTUAL_ENV only and let the
# activation block at the end of .zshrc order PATH correctly. Pre-activating here
# would leave venv/bin behind gnubin and homebrew, same as tmux-inherited venvs.
codex() {
    if [[ -n ${VIRTUAL_ENV:-} ]]; then
        command codex "$@"
    else
        VIRTUAL_ENV=$(_agent_venv_path) command codex "$@"
    fi
}