#!/usr/bin/env bash
# pyreplab — CLI for the persistent Python REPL
#
# Usage:
#   pyreplab start [--workdir DIR] [--cwd DIR] [--venv PATH] [...]
#   pyreplab run notebook.py       Run all cells (stamps [N] into # %% markers)
#   pyreplab run notebook.py:2     Run cell 2
#   pyreplab run 'print("hello")'  Inline code
#   pyreplab run < script.py       Read from stdin
#   pyreplab cells notebook.py     List cells (stamps [N], peeks comments)
#   pyreplab wait                  Wait for a running command to finish
#   pyreplab cancel                Cancel the currently running command
#   pyreplab status                Check if REPL is running
#   pyreplab ps                    List all active sessions
#   pyreplab stop / stop-all       Stop session(s)
#   pyreplab clean                 Remove session files
#
# No sourcing needed. Each invocation is standalone.

set -euo pipefail

PYREPLAB_BASE="${PYREPLAB_BASE:-/tmp/pyreplab}"
# Resolve symlinks (uv tool installs symlink this script into ~/.local/bin)
# so the real location inside the tool venv is found
_self="$0"
while [ -L "$_self" ]; do
    _link="$(readlink "$_self")"
    case "$_link" in
        /*) _self="$_link" ;;
        *) _self="$(cd "$(dirname "$_self")" && pwd)/$_link" ;;
    esac
done
_self_dir="$(cd "$(dirname "$_self")" && pwd)"
# Prefer the python next to this script (uv tool installs place both the
# script and the venv python in bin/); fall back to PATH python3 for
# repo/dev installs.
if [ -x "$_self_dir/python" ]; then
    PYREPLAB_PYTHON="$_self_dir/python"
else
    PYREPLAB_PYTHON="python3"
fi
if [ -n "${PYREPLAB_SCRIPT:-}" ]; then
    :
elif [ -f "$_self_dir/pyreplab.py" ]; then
    PYREPLAB_SCRIPT="$_self_dir/pyreplab.py"
else
    # uv tool installs put pyreplab.py in the venv's site-packages: the
    # venv's own python (adjacent to this script) can import it, while a
    # PATH python3 usually cannot
    PYREPLAB_SCRIPT="$("$PYREPLAB_PYTHON" -c 'import pyreplab; print(pyreplab.__file__)' 2>/dev/null)" || PYREPLAB_SCRIPT=""
fi

# Discover the project root for a directory: nearest ancestor with .git,
# else nearest with pyproject.toml, else the directory itself.
_discover_root() {
    local dir
    dir=$(cd -P "$1" 2>/dev/null && pwd -P) || dir=$(pwd -P 2>/dev/null || echo /)
    local cur="$dir"
    while :; do
        if [ -e "$cur/.git" ]; then
            echo "$cur"
            return 0
        fi
        if [ -f "$cur/pyproject.toml" ]; then
            echo "$cur"
            return 0
        fi
        [ "$cur" = "/" ] && break
        cur=$(dirname "$cur")
    done
    echo "$dir"
}

# Resolve session dir: explicit PYREPLAB_DIR, or derived from the project root
_resolve_dir() {
    if [ -n "${PYREPLAB_DIR:-}" ]; then
        echo "$PYREPLAB_DIR"
    else
        local start="${1:-${_PYREPLAB_WORKDIR:-$(pwd -P 2>/dev/null || pwd)}}"
        local root
        if [ "${_PYREPLAB_WORKDIR_EXPLICIT:-0}" = "1" ] && [ -z "$1" ]; then
            # Legacy --workdir on stop/clean/dir etc. is an exact root
            root="$start"
        else
            root=$(_discover_root "$start")
        fi
        _session_dir_for_root "$root"
    fi
}

# Hash a project root into a session directory name
_session_dir_for_root() {
    local root="$1" hash name
    hash=$(printf '%s' "$root" | md5sum 2>/dev/null | cut -c1-8 || printf '%s' "$root" | md5 -q 2>/dev/null | cut -c1-8)
    name=$(basename "$root")
    echo "$PYREPLAB_BASE/${name}_${hash}"
}

# Extract --workdir from args, or default to cwd. Explicit --workdir is a
# legacy exact root (used by stop/clean/dir etc.); otherwise the caller's
# cwd is the discovery starting point.
_PYREPLAB_WORKDIR=""
_PYREPLAB_WORKDIR_EXPLICIT=0
_parse_workdir() {
    local args=("$@")
    for ((i=0; i<${#args[@]}; i++)); do
        if [ "${args[$i]}" = "--workdir" ] && [ $((i+1)) -lt ${#args[@]} ]; then
            _PYREPLAB_WORKDIR="$(cd "${args[$((i+1))]}" && pwd -P)"
            _PYREPLAB_WORKDIR_EXPLICIT=1
            return
        fi
    done
    # Default: cwd
    _PYREPLAB_WORKDIR="$(pwd -P 2>/dev/null || pwd)"
    _PYREPLAB_WORKDIR_EXPLICIT=0
}

_is_running() {
    local dir
    dir=$(_resolve_dir)
    local pidfile="$dir/pyreplab.pid"
    [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile")" 2>/dev/null
}

# Start a session: resolves --workdir (deprecated), --cwd and --session-dir
# (deprecated) once, so the daemon never sees conflicting path flags.
cmd_start() {
    # Parse the path-bearing flags; everything else is passed through
    local workdir="" cwd="" session_dir=""
    local args=("$@") rest=() i=0
    while [ $i -lt ${#args[@]} ]; do
        case "${args[$i]}" in
            --workdir)
                if [ $((i+1)) -lt ${#args[@]} ]; then workdir="${args[$((i+1))]}"; fi
                i=$((i+2)) ;;
            --cwd)
                if [ $((i+1)) -lt ${#args[@]} ]; then cwd="${args[$((i+1))]}"; fi
                i=$((i+2)) ;;
            --session-dir)
                if [ $((i+1)) -lt ${#args[@]} ]; then session_dir="${args[$((i+1))]}"; fi
                i=$((i+2)) ;;
            *)
                rest+=("${args[$i]}")
                i=$((i+1)) ;;
        esac
    done

    if [ -n "$workdir" ]; then
        echo "pyreplab: --workdir is deprecated: sessions are auto-discovered from the project root (nearest .git or pyproject.toml); pass --cwd to pin the execution directory" >&2
        workdir=$(cd "$workdir" && pwd -P)
    fi
    if [ -n "$session_dir" ]; then
        echo "pyreplab: --session-dir is deprecated: use PYREPLAB_DIR instead" >&2
    fi
    if [ -n "$cwd" ]; then
        cwd=$(cd -P "$cwd" && pwd -P) || {
            echo "pyreplab: --cwd not found: $cwd" >&2
            return 1
        }
    fi

    # Session identity: legacy --workdir wins, else discovery from --cwd (or
    # the start directory). Session dir: --session-dir > PYREPLAB_DIR > hash.
    local root dir
    if [ -n "$workdir" ]; then
        root="$workdir"
    else
        root=$(_discover_root "${cwd:-$(pwd -P 2>/dev/null || pwd)}")
    fi
    if [ -n "$session_dir" ]; then
        dir="$session_dir"
    elif [ -n "${PYREPLAB_DIR:-}" ]; then
        dir="$PYREPLAB_DIR"
    else
        dir=$(_session_dir_for_root "$root")
    fi

    local pidfile="$dir/pyreplab.pid"

    if [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile")" 2>/dev/null; then
        echo "pyreplab: already running (pid $(cat "$pidfile"), dir $dir)" >&2
        return 0
    fi
    mkdir -p "$dir"
    # Clean stale IPC files from a previous (dead) session
    rm -f "$dir/cmd.py" "$dir/output.json" "$dir/done" "$dir/pending_id" "$dir/pending_start" "$dir/progress.json" "$dir/progress.json.tmp"
    # Detach the daemon from this shell session: ignore SIGHUP (nohup), drop
    # the session's stdin/stdout/stderr (log to daemon.log), and remove it from
    # the job table (disown). The daemon survives the bash session ending and
    # never holds the caller's pipes open (which would hang output capture).
    local daemon_args=(--session-dir "$dir" --session-root "$root")
    if [ -n "$cwd" ]; then
        daemon_args+=(--cwd "$cwd")
    fi
    daemon_args+=("${rest[@]}")
    if [ -z "$PYREPLAB_SCRIPT" ]; then
        echo "pyreplab: cannot locate pyreplab.py — reinstall, or set PYREPLAB_SCRIPT" >&2
        return 1
    fi
    nohup "$PYREPLAB_PYTHON" "$PYREPLAB_SCRIPT" "${daemon_args[@]}" >"$dir/daemon.log" 2>&1 </dev/null &
    local pid=$!
    disown "$pid" 2>/dev/null || true
    echo "$pid" > "$pidfile"
    sleep 0.3
    if ! kill -0 "$pid" 2>/dev/null; then
        echo "pyreplab: failed to start" >&2
        rm -f "$pidfile"
        return 1
    fi
    echo "pyreplab: started (pid $pid, dir $dir)" >&2
}

cmd_stop() {
    local dir
    dir=$(_resolve_dir)
    local pidfile="$dir/pyreplab.pid"
    if ! [ -f "$pidfile" ] || ! kill -0 "$(cat "$pidfile")" 2>/dev/null; then
        echo "pyreplab: not running" >&2
        rm -f "$pidfile"
        return 0
    fi
    local pid
    pid=$(cat "$pidfile")
    kill "$pid" 2>/dev/null
    local i=0
    while kill -0 "$pid" 2>/dev/null && [ "$i" -lt 30 ]; do
        sleep 0.1
        i=$((i + 1))
    done
    rm -f "$pidfile"
    echo "pyreplab: stopped (pid $pid)" >&2
}

cmd_clean() {
    local dir
    dir=$(_resolve_dir)
    for f in cmd.py cmd.py.tmp output.json output.json.tmp done pending_id pending_start progress.json progress.json.tmp; do
        [ -f "$dir/$f" ] && rm -f "$dir/$f"
    done
    echo "pyreplab: cleaned session files in $dir" >&2
}

_count_cells() {
    # Count the number of cells in a .py file.
    local file="$1"
    python3 -c '
import sys, re
text = open(sys.argv[1]).read()
markers = list(re.finditer(r"(?m)^# ?%%[^\n]*\n", text))
if not markers:
    print(1)
elif re.match(r"# ?%%", text):
    print(len(markers))
else:
    print(len(markers) + 1)
' "$file"
}

_stamp_cells() {
    # Add or update [N] indices on cell markers in a .py file. Idempotent.
    local file="$1"
    python3 -c '
import sys, re

path = sys.argv[1]
text = open(path).read()
lines = text.split("\n")

# If file starts with a cell marker, first marker is cell 0.
# Otherwise preamble is cell 0, first marker is cell 1.
has_preamble = not re.match(r"# ?%%", text)
cell_idx = 1 if has_preamble else 0
changed = False

for i, line in enumerate(lines):
    m = re.match(r"^(# ?%%)\s*(?:\[(\d+)\]\s*)?(.*?)$", line)
    if m:
        prefix, existing_idx, label = m.groups()
        new_line = f"# %% [{cell_idx}]"
        if label.strip():
            new_line += f" {label.strip()}"
        if lines[i] != new_line:
            lines[i] = new_line
            changed = True
        cell_idx += 1

if changed:
    with open(path, "w") as f:
        f.write("\n".join(lines))
' "$file"
}

_extract_cell() {
    # Parse a .py file into #%% cells and extract one by index.
    # Usage: _extract_cell file.py [cell_index]
    # If no index, returns the whole file.
    local file="$1"
    local cell_idx="${2:-all}"

    if [ "$cell_idx" = "all" ]; then
        cat "$file"
        return
    fi

    python3 -c '
import sys, re

text = open(sys.argv[1]).read()
# Split on # %% or #%% lines (accept optional space per industry convention)
parts = re.split(r"(?m)^# ?%%[^\n]*\n", text)

# If file starts with a cell marker, first split is empty string before it
# If file does NOT start with one, first split is the preamble (cell 0)
cells = []
markers = list(re.finditer(r"(?m)^# ?%%[^\n]*\n", text))

if re.match(r"# ?%%", text):
    # Each marker corresponds to a cell
    cells = parts[1:]  # skip empty first split
else:
    # parts[0] is preamble before first marker, rest follow markers
    cells = parts

idx = int(sys.argv[2])
if idx < 0 or idx >= len(cells):
    print(f"pyreplab: cell {idx} not found (file has {len(cells)} cells)", file=sys.stderr)
    sys.exit(1)
print(cells[idx], end="")
' "$file" "$cell_idx"
}

# Print a progress line to stderr when progress.json has new data from the
# daemon (partial output during long-running commands). Dedup: once per new
# snapshot for chatty commands; heartbeat every 10s for silent ones.
# State is per-invocation, so a fresh `wait` shows the current progress once.
_PROG_FKEY=""
_PROG_CKEY=""
_PROG_TS=""
_show_progress() {
    [ "${PYREPLAB_PROGRESS:-1}" = "1" ] || return 0
    local dir="$1" cmd_id="${2:-}" p="$dir/progress.json"
    [ -f "$p" ] || return 0
    # Cheap pre-filter: only parse when the file changed (mtime+size — size
    # catches sub-second snapshots where mtime resolution is 1s)
    local pm="" pz=""
    pm=$(stat -f%m "$p" 2>/dev/null) || pm=$(stat -c%Y "$p" 2>/dev/null) || true
    pz=$(stat -f%z "$p" 2>/dev/null) || pz=$(stat -c%s "$p" 2>/dev/null) || true
    [ -n "$pm" ] || return 0
    local fkey="${pm}:${pz}"
    [ "$fkey" = "$_PROG_FKEY" ] && return 0
    _PROG_FKEY="$fkey"
    local content=""
    content=$(cat "$p" 2>/dev/null) || true
    [ -z "$content" ] && return 0
    local kv=""
    kv=$(printf '%s' "$content" | python3 -c '
import json, sys
d = json.load(sys.stdin)
want = sys.argv[1] if len(sys.argv) > 1 else ""
if want and d.get("id") != want:
    sys.exit(0)
chars = d.get("stdout_chars", 0) + d.get("stderr_chars", 0)
cell = d.get("cell", "")
elapsed = d.get("elapsed", 0)
tail = (d.get("stdout") or "").rstrip().splitlines()
if not tail:
    tail = (d.get("stderr") or "").rstrip().splitlines()
last = tail[-1].strip()[-120:] if tail else ""
parts = []
if cell: parts.append("cell %s" % cell)
parts.append("%d chars out" % chars)
msg = "pyreplab: progress (%.1fs) %s" % (elapsed, " | ".join(parts))
if last: msg += " | last: %s" % last
print("%d|%s" % (chars, cell))
print(msg)
' "$cmd_id" 2>/dev/null) || kv=""
    [ -z "$kv" ] && return 0
    local ckey="${kv%%$'\n'*}"
    local msg="${kv#*$'\n'}"
    if [ "$ckey" = "$_PROG_CKEY" ]; then
        # Same output as before: heartbeat only every 10s
        local now
        now=$(date +%s)
        [ -n "$_PROG_TS" ] && [ $((now - _PROG_TS)) -lt 10 ] && return 0
        _PROG_TS="$now"
    else
        _PROG_TS=$(date +%s)
    fi
    _PROG_CKEY="$ckey"
    echo "$msg" >&2
}

_wait_for_result() {
    # Poll for command completion and print output.
    # Returns 0 on success, 1 on error output, 2 on timeout (still running).
    local timeout="${1:-30}"
    local cmd_id="${2:-}"
    local dir
    dir=$(_resolve_dir)
    local done_path="$dir/done"
    local output_path="$dir/output.json"
    local pending_path="$dir/pending_id"

    # Read submission timestamp for total elapsed reporting
    local start_time=""
    [ -f "$dir/pending_start" ] && start_time=$(cat "$dir/pending_start")

    # Poll for completion. If a done/output.json pair appears whose id does not
    # match this command (stale files from a previous command or an orphaned
    # daemon), discard it and keep polling for the real result.
    local elapsed=0
    local result=""
    while :; do
        while [ ! -f "$done_path" ]; do
            sleep 0.1
            elapsed=$((elapsed + 1))
            _show_progress "$dir" "$cmd_id"
            if [ "$elapsed" -ge "$((timeout * 10))" ]; then
                local total="${timeout}s"
                if [ -n "$start_time" ]; then
                    total="$(( $(date +%s) - start_time ))s"
                fi
                echo "pyreplab: still running (${total} elapsed). Run \`pyreplab wait\` to check again." >&2
                return 2
            fi
        done

        # Read output and verify it belongs to this command
        result=$(cat "$output_path" 2>/dev/null) || true
        if [ -z "$result" ]; then
            echo "pyreplab: no output received" >&2
            rm -f "$done_path" "$pending_path" "$dir/pending_start" "$dir/progress.json" "$dir/progress.json.tmp"
            return 1
        fi
        if [ -n "$cmd_id" ]; then
            local result_id=""
            result_id=$(printf '%s' "$result" | python3 -c 'import json,sys; print(json.load(sys.stdin).get("id",""))' 2>/dev/null) || true
            if [ -n "$result_id" ] && [ "$result_id" != "$cmd_id" ]; then
                echo "pyreplab: discarding stale output (id $result_id != $cmd_id), waiting for the real result..." >&2
                rm -f "$done_path" "$output_path"
                continue
            fi
        fi
        break
    done

    # Clean up handshake files
    rm -f "$done_path"
    rm -f "$output_path"
    rm -f "$pending_path" "$dir/pending_start" "$dir/progress.json" "$dir/progress.json.tmp"

    # Parse JSON and print (single python3 call)
    printf '%s' "$result" | python3 -c '
import json, sys
d = json.load(sys.stdin)
stdout = d.get("stdout", "")
stderr = d.get("stderr", "")
error = d.get("error")
if stdout:
    print(stdout, end="")
if stderr:
    print(stderr, end="", file=sys.stderr)
if error:
    print(error, end="", file=sys.stderr)
    sys.exit(1)
'
}

# Verify a session is usable before sending a command: the session dir must
# exist and its daemon must be alive.
_session_ready() {
    local dir="$1"
    if [ ! -d "$dir" ]; then
        echo "pyreplab: no session for this directory — run \`pyreplab start\` first (or set PYREPLAB_DIR)" >&2
        return 1
    fi
    if ! _is_running; then
        echo "pyreplab: session not running — run \`pyreplab start\` first (or set PYREPLAB_DIR)" >&2
        return 1
    fi
    return 0
}

_send_code() {
    # Send code to the pyreplab server and print the response.
    local code="$1"
    local cell_label="${2:-}"
    local dir
    dir=$(_resolve_dir)
    if ! _session_ready "$dir"; then
        return 1
    fi
    local id
    id="cmd_$$_$(date +%s%N 2>/dev/null || echo $RANDOM)"
    local cmd_path="$dir/cmd.py"
    local done_path="$dir/done"
    local output_path="$dir/output.json"
    local pending_path="$dir/pending_id"
    local timeout="${PYREPLAB_TIMEOUT:-30}"

    # Check if server is busy with a previous command
    if [ -f "$pending_path" ]; then
        # If daemon is dead, these are stale files from a previous session — clean up
        if ! _is_running; then
            rm -f "$pending_path" "$done_path" "$output_path" "$dir/cmd.py" "$dir/pending_start" "$dir/progress.json" "$dir/progress.json.tmp"
        else
            echo "pyreplab: busy running previous command. Run \`pyreplab wait\` first." >&2
            return 1
        fi
    fi

    # Remove stale handshake files
    rm -f "$done_path" "$output_path" "$dir/progress.json" "$dir/progress.json.tmp"

    # Write cmd.py with #%% cell header (atomic: write tmp, then move)
    local header="id: $id cwd: $(pwd)"
    [ -n "$cell_label" ] && header="$header cell: $cell_label"
    printf '#%%%% %s\n%s\n' "$header" "$code" > "$cmd_path.tmp"
    mv "$cmd_path.tmp" "$cmd_path"

    # Track pending command and start time
    printf '%s' "$id" > "$pending_path"
    date +%s > "$dir/pending_start"

    # Wait for result
    _wait_for_result "$timeout" "$id"
}

_send_notebook() {
    # Send a notebook path to the daemon for server-side multi-cell execution.
    # The daemon reads the file, splits cells, and runs them all sequentially.
    local notebook_abs="$1"
    local dir
    dir=$(_resolve_dir)
    local id
    id="cmd_$$_$(date +%s%N 2>/dev/null || echo $RANDOM)"
    local cmd_path="$dir/cmd.py"
    local done_path="$dir/done"
    local output_path="$dir/output.json"
    local pending_path="$dir/pending_id"
    local timeout="${PYREPLAB_TIMEOUT:-30}"

    # Check if server is busy with a previous command
    if [ -f "$pending_path" ]; then
        if ! _is_running; then
            rm -f "$pending_path" "$done_path" "$output_path" "$dir/cmd.py" "$dir/pending_start" "$dir/progress.json" "$dir/progress.json.tmp"
        else
            echo "pyreplab: busy running previous command. Run \`pyreplab wait\` first." >&2
            return 1
        fi
    fi

    rm -f "$done_path" "$output_path" "$dir/progress.json" "$dir/progress.json.tmp"

    # Write cmd.py with notebook header (no code body — daemon reads the file)
    local header="id: $id cwd: $(pwd) notebook: $notebook_abs"
    printf '#%%%% %s\n' "$header" > "$cmd_path.tmp"
    mv "$cmd_path.tmp" "$cmd_path"

    printf '%s' "$id" > "$pending_path"
    date +%s > "$dir/pending_start"

    _wait_for_result "$timeout"
}

cmd_cells() {
    local file="${1:-}"
    if [ -z "$file" ] || ! [ -f "$file" ]; then
        echo "Usage: pyreplab cells <file.py>" >&2
        return 1
    fi
    local stamp="${PYREPLAB_STAMP:-1}"
    [ "$stamp" = "1" ] && _stamp_cells "$file"
    python3 -c '
import sys, re
text = open(sys.argv[1]).read()
lines = text.split("\n")
markers = list(re.finditer(r"(?m)^# ?%%[^\n]*$", text))
if not markers:
    # No cell markers — whole file is cell 0
    first = lines[0].strip() if lines else "(empty)"
    print(f"  0: {first}")
    sys.exit(0)
# If file starts with content before first marker, that is cell 0 (preamble)
if markers[0].start() > 0:
    pre = text[:markers[0].start()].strip().split("\n")
    first = next((l.strip() for l in pre if l.strip()), "(empty)")
    print(f"  0: {first}")
    offset = 1
else:
    offset = 0
for i, m in enumerate(markers):
    # Strip the marker prefix to get the inline label
    label = re.sub(r"^# ?%%\s*(?:\[\d+\]\s*)?", "", m.group()).strip()
    # If no inline label, peek at the next line for a comment
    if not label:
        end = m.end()
        line_start = text.find("\n", end)
        if line_start == -1:
            next_line = ""
        else:
            # The next line starts right after the marker line
            next_start = end + 1 if end < len(text) and text[end] == "\n" else end
            next_end = text.find("\n", next_start)
            next_line = text[next_start:next_end].strip() if next_end != -1 else text[next_start:].strip()
        if next_line.startswith("#") and not re.match(r"# ?%%", next_line):
            label = next_line.lstrip("# ").strip()
    label = label or "(unnamed)"
    print(f"  {i + offset}: # %% {label}")
' "$file"
}

cmd_run() {
    # Resolve session from cwd (unless PYREPLAB_DIR is set)
    [ -z "${PYREPLAB_DIR:-}" ] && _parse_workdir
    local arg="${1:-}"
    local stamp="${PYREPLAB_STAMP:-1}"

    if [ -z "$arg" ]; then
        # No args: read from stdin
        local code
        code=$(cat)
        if [ -z "$code" ]; then
            echo "pyreplab: no code to run" >&2
            return 1
        fi
        _send_code "$code" ""
    elif [ -f "$arg" ]; then
        # pyreplab run file.py — stamp cells, send notebook to daemon for server-side execution
        [ "$stamp" = "1" ] && _stamp_cells "$arg"
        local file_abs
        file_abs="$(cd "$(dirname "$arg")" && pwd)/$(basename "$arg")"
        _send_notebook "$file_abs"
    elif [[ "$arg" == *:* ]] && [ -f "${arg%%:*}" ]; then
        # pyreplab run file.py:N — stamp cells, then run cell N
        local file="${arg%%:*}"
        local cell="${arg##*:}"
        [ "$stamp" = "1" ] && _stamp_cells "$file"
        local code
        code=$(_extract_cell "$file" "$cell")
        if [ -z "$code" ]; then
            echo "pyreplab: no code to run" >&2
            return 1
        fi
        _send_code "$code" "$(basename "$file"):$cell"
    else
        # Inline code
        _send_code "$arg" ""
    fi
}

cmd_wait() {
    local dir
    dir=$(_resolve_dir)

    if [ ! -f "$dir/pending_id" ]; then
        echo "pyreplab: no command pending" >&2
        return 1
    fi

    # If daemon is dead, the pending command will never complete — clean up
    if ! _is_running; then
        rm -f "$dir/pending_id" "$dir/pending_start" "$dir/done" "$dir/output.json" "$dir/cmd.py" "$dir/progress.json" "$dir/progress.json.tmp"
        echo "pyreplab: server died while command was pending (cleaned up stale files)" >&2
        return 1
    fi

    # Short poll (2s) — return quickly so agents aren't blocked
    _wait_for_result 2 "$(cat "$dir/pending_id" 2>/dev/null || true)"
}

cmd_cancel() {
    local dir
    dir=$(_resolve_dir)
    local pidfile="$dir/pyreplab.pid"

    if ! [ -f "$pidfile" ] || ! kill -0 "$(cat "$pidfile")" 2>/dev/null; then
        echo "pyreplab: not running" >&2
        return 1
    fi

    if [ ! -f "$dir/pending_id" ]; then
        echo "pyreplab: no command running" >&2
        return 0
    fi

    local pid
    pid=$(cat "$pidfile")
    kill -USR1 "$pid" 2>/dev/null
    echo "pyreplab: cancel signal sent" >&2

    # Wait briefly for the daemon to finish processing the cancellation
    _wait_for_result 5 "$(cat "$dir/pending_id" 2>/dev/null || true)"
}

cmd_dir() {
    local dir
    dir=$(_resolve_dir)
    echo "$dir"
}

cmd_status() {
    local dir
    dir=$(_resolve_dir)
    local pidfile="$dir/pyreplab.pid"
    if [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile")" 2>/dev/null; then
        local state="idle"
        if [ -f "$dir/pending_id" ]; then
            state="executing command"
            if [ -f "$dir/pending_start" ]; then
                local elapsed=$(( $(date +%s) - $(cat "$dir/pending_start") ))
                state="executing command (${elapsed}s elapsed)"
            fi
            # Append latest progress tail from the daemon's progress.json
            # (only if it belongs to the pending command)
            if [ -f "$dir/progress.json" ]; then
                local ptail="" want_id=""
                want_id=$(cat "$dir/pending_id" 2>/dev/null) || true
                ptail=$(python3 -c '
import json, sys
d = json.load(open(sys.argv[1]))
want = sys.argv[2] if len(sys.argv) > 2 else ""
if want and d.get("id") != want:
    sys.exit(0)
chars = d.get("stdout_chars", 0) + d.get("stderr_chars", 0)
tail = (d.get("stdout") or "").rstrip().splitlines()
if not tail:
    tail = (d.get("stderr") or "").rstrip().splitlines()
last = tail[-1].strip()[-80:] if tail else ""
s = "progress %.1fs | %d chars" % (d.get("elapsed", 0), chars)
if d.get("cell"): s += " | cell %s" % d["cell"]
if last: s += " | last: %s" % last
print(s)
' "$dir/progress.json" "$want_id" 2>/dev/null) || ptail=""
                [ -n "$ptail" ] && state="$state, $ptail"
            fi
        fi
        echo "pyreplab: running (pid $(cat "$pidfile"), dir $dir), $state"
        # Show resolved session configuration (root, mode, env, python)
        if [ -f "$dir/session.json" ]; then
            local sjson=""
            sjson=$(python3 -c '
import json, sys
d = json.load(open(sys.argv[1]))
env = d.get("env", {})
env_s = "none"
if env.get("path"):
    env_s = "%s:%s" % (env.get("type", "?"), env.get("path", "?"))
s = "root %s | mode %s | env %s | python %s" % (
    d.get("session_root", "?"), d.get("mode", "?"), env_s,
    env.get("version", "?"))
print(s)
' "$dir/session.json" 2>/dev/null) || sjson=""
            [ -n "$sjson" ] && echo "pyreplab: $sjson" >&2
        fi
    else
        echo "pyreplab: not running"
        return 1
    fi
}

cmd_ps() {
    local found=0
    local format="%-28s %-7s %-8s %-6s %s\n"
    printf "$format" "SESSION" "PID" "UPTIME" "MEM" "DIR"
    for pidfile in "$PYREPLAB_BASE"/*/pyreplab.pid; do
        [ -f "$pidfile" ] || continue
        local pid dir name
        pid=$(cat "$pidfile")
        dir=$(dirname "$pidfile")
        name=$(basename "$dir")
        if kill -0 "$pid" 2>/dev/null; then
            # Uptime from pidfile creation time
            local uptime="?"
            if stat -f%m "$pidfile" >/dev/null 2>&1; then
                # macOS stat
                local created now elapsed
                created=$(stat -f%m "$pidfile")
                now=$(date +%s)
                elapsed=$((now - created))
            elif stat -c%Y "$pidfile" >/dev/null 2>&1; then
                # Linux stat
                local created now elapsed
                created=$(stat -c%Y "$pidfile")
                now=$(date +%s)
                elapsed=$((now - created))
            fi
            if [ -n "${elapsed:-}" ]; then
                if [ "$elapsed" -ge 3600 ]; then
                    uptime="$((elapsed / 3600))h$((elapsed % 3600 / 60))m"
                elif [ "$elapsed" -ge 60 ]; then
                    uptime="$((elapsed / 60))m$((elapsed % 60))s"
                else
                    uptime="${elapsed}s"
                fi
            fi
            # Memory (RSS in MB)
            local mem="?"
            local rss
            rss=$(ps -o rss= -p "$pid" 2>/dev/null | tr -d ' ')
            if [ -n "$rss" ]; then
                mem="$((rss / 1024))MB"
            fi
            printf "$format" "$name" "$pid" "$uptime" "$mem" "$dir"
            found=$((found + 1))
        else
            rm -f "$pidfile"
        fi
    done
    if [ "$found" -eq 0 ]; then
        echo "pyreplab: no active sessions"
    fi
}

cmd_stop_all() {
    local stopped=0
    for pidfile in "$PYREPLAB_BASE"/*/pyreplab.pid; do
        [ -f "$pidfile" ] || continue
        local pid dir name
        pid=$(cat "$pidfile")
        dir=$(dirname "$pidfile")
        name=$(basename "$dir")
        if kill -0 "$pid" 2>/dev/null; then
            kill "$pid" 2>/dev/null
            local i=0
            while kill -0 "$pid" 2>/dev/null && [ "$i" -lt 30 ]; do
                sleep 0.1
                i=$((i + 1))
            done
            echo "pyreplab: stopped $name (pid $pid)" >&2
            stopped=$((stopped + 1))
        fi
        rm -f "$pidfile"
    done
    if [ "$stopped" -eq 0 ]; then
        echo "pyreplab: no sessions to stop" >&2
    else
        echo "pyreplab: stopped $stopped session(s)" >&2
    fi
}

# --- Main dispatch ---
case "${1:-help}" in
    start)  shift; cmd_start "$@" ;;
    stop)   shift; _parse_workdir "$@"; cmd_stop ;;
    stop-all) cmd_stop_all ;;
    clean)  shift; _parse_workdir "$@"; cmd_clean ;;
    run)    shift; cmd_run "$@" ;;
    wait)   shift; _parse_workdir "$@"; cmd_wait ;;
    cancel) shift; _parse_workdir "$@"; cmd_cancel ;;
    cells)  shift; cmd_cells "$@" ;;
    dir)    shift; _parse_workdir "$@"; cmd_dir ;;
    status) shift; _parse_workdir "$@"; cmd_status ;;
    ps|list|ls) cmd_ps ;;
    help|--help|-h)
        echo "Usage: pyreplab <command> [args]"
        echo ""
        echo "Commands:"
        echo "  start [opts]        Start the REPL (see START OPTIONS below)"
        echo "  run file.py         Run all cells (stamps [N] indices into file)"
        echo "  run file.py:N       Run cell N from file (0-indexed)"
        echo "  run 'code'          Run inline code"
        echo "  run                 Read code from stdin"
        echo "  cells file.py       List cells (stamps [N] indices into file)"
        echo "  wait                Wait for a running command to finish"
        echo "  cancel              Cancel the currently running command"
        echo "  stop                Stop the current session"
        echo "  stop-all            Stop all active sessions"
        echo "  dir                 Print session directory path"
        echo "  status              Check if REPL is running"
        echo "  ps                  List all active sessions"
        echo "  clean               Remove session files"
        echo ""
        echo "AGENT WORKFLOW (how to drive pyreplab)"
        echo "  # 1. Ensure a session is running for this directory:"
        echo "  pyreplab status || pyreplab start"
        echo "  # 2. Execute code: a file (all cells), one cell, or inline:"
        echo "  pyreplab run analysis.py | pyreplab run analysis.py:2 | pyreplab run 'print(x)'"
        echo "  # 3. run exits: 0 = done (output printed), 1 = error, 2 = still running."
        echo "  #    On exit 2, poll until it finishes:"
        echo "  pyreplab wait    # repeat until exit 0 or 1 (2s poll each)"
        echo "  pyreplab cancel  # interrupt the running command (session survives)"
        echo "  # 4. Long loops stream progress to stderr during run/wait:"
        echo "  #    pyreplab: progress (3.0s) 710 chars out | last: iter 89"
        echo "  # 5. Everything persists: imports, variables and state survive"
        echo "  #    between commands and cells (one session per project)."
        echo ""
        echo "START OPTIONS (pyreplab start)"
        echo "  --cwd DIR           Lock the REPL's working directory to DIR (sticky)."
        echo "                      Default: each command runs in the caller's shell"
        echo "                      directory, so relative imports and file paths"
        echo "                      resolve where you are."
        echo "  --workdir DIR       [DEPRECATED] project root for session identity;"
        echo "                      sessions are now auto-discovered instead."
        echo "  --venv PATH         Explicit virtualenv (venv/uv) to activate."
        echo "  --conda [ENV]       Activate a conda environment (default: base)."
        echo "  --no-conda          Disable automatic conda base fallback."
        echo "  --max-output N      Max output chars per command (default: 100000)."
        echo "  --max-rows N        Pandas max display rows (default: 50)."
        echo "  --max-cols N        Pandas max display columns (default: 20)."
        echo "  --poll-interval S   Daemon cmd.py poll interval (default: 0.05)."
        echo "  --progress-interval S  Progress snapshot interval, 0 disables (default: 1.0)."
        echo ""
        echo "ENVIRONMENT AUTO-DETECTION (agents don't need to know the env system)"
        echo "  Priority: --venv > \$PIXI_ENVIRONMENT_PREFIX (pixi run/shell) >"
        echo "  nearest .venv (venv/uv) or .pixi/envs/<name> (pixi) between the"
        echo "  working directory and the project root > conda base fallback."
        echo "  The daemon re-executes under the env's own python, so versions"
        echo "  always match. Disable the conda fallback with --no-conda."
        echo ""
        echo "SESSIONS"
        echo "  One session per project root: auto-discovered from the nearest"
        echo "  ancestor containing .git or pyproject.toml (or the start dir)."
        echo "  Override the session directory with PYREPLAB_DIR."
        echo ""
        echo "ENVIRONMENT VARIABLES"
        echo "  PYREPLAB_DIR        Session directory (overrides discovery)"
        echo "  PYREPLAB_TIMEOUT    Client poll timeout in seconds (default: 30)"
        echo "  PYREPLAB_PROGRESS   Set to 0 to disable progress lines"
        echo "  PYREPLAB_STAMP      Set to 0 to disable [N] cell stamping"
        echo ""
        echo "Cell stamping adds [N] indices to # %% markers in your .py files."
        echo "Accepts both # %% and #%% markers."
        ;;
    *)
        echo "pyreplab: unknown command '$1' (try 'pyreplab help')" >&2
        exit 1
        ;;
esac
