#!/usr/bin/env bash
# claude-sandbox — helper CLI, placed on PATH at /usr/local/bin by
# install.sh. The install clone is disposable: every command here works
# with no clone present (update fetches its own fresh one).
set -euo pipefail

REPO_URL="https://github.com/DiamondLightSource/claude-sandbox"
# Stamped by install.sh from `git describe` on the installing clone.
# Seam for tests only (version reporting is not security-critical).
VERSION_FILE="${CLAUDE_SANDBOX_VERSION_FILE:-/usr/libexec/claude-sandbox/version}"
# Written by install.sh only when the installer was the PyPI wheel.
INSTALLER_FILE="${CLAUDE_SANDBOX_INSTALLER_FILE:-/usr/libexec/claude-sandbox/installer}"

usage() {
    cat <<'EOF'
claude-sandbox — sandboxed-Claude helper commands

Usage: claude-sandbox <command> [args]

Commands:
  install               Explain that the sandbox is already installed
  gh-auth               Authenticate gh with a GitHub PAT (kept out of shell history)
  glab-auth [hostname]  Authenticate glab with a GitLab PAT (default hostname: gitlab.diamond.ac.uk)
  update                Clone and install the latest claude-sandbox release
  verify [--agent NAME] Run the isolation battery directly (no agent login needed)
                        (NAME: claude, the default, codex, or pi)
  pi-local [--port PORT] Discover the running lllm2 model and context (port 1920)
  pi-local MODEL CONTEXT [PORT]
                        Configure Pi's lllm2 provider manually
  doctor [--fix]        Check the recommended setup: container tag in the Claude status
                        line, Pi footer and shell prompts. --fix applies it, with backups
  version               Show the installed claude-sandbox version
  help                  Show this help
EOF
}

# Authenticate gh CLI with a GitHub PAT (token not stored in shell history).
cmd_gh_auth() {
    local url t
    url=$'\e[4;36mhttps://github.com/settings/personal-access-tokens\e[0m'
    cat <<EOF
Create or renew a fine-grained PAT at:
  $url

Recommended settings for a sandboxed Claude Code:
  - Resource owner: your user (or org that owns this repo)
  - Repository access: Only select repositories -> just this repo
  - Expiration: short (e.g. 30 days) so a leaked token expires quickly
  - Repository permissions Read/Write:
      Issues, Pull requests
  - Repository permissions Read Only:
      Contents
    (Metadata: Read-only is added automatically)
  - Leave everything else unset / no access

EOF
    read -sp "GitHub PAT: " t && echo
    echo "$t" | gh auth login --with-token
    unset t
    gh auth setup-git
    gh auth status
}

# Authenticate glab CLI with a GitLab PAT (token not stored in shell history)
# and pin the instance to HTTPS, so glab never drives git over SSH.
cmd_glab_auth() {
    local hostname="${1:-gitlab.diamond.ac.uk}" url t
    # Derive the PAT page from the instance being authenticated. A hardcoded
    # gitlab.com URL sends self-hosted users to the wrong site to mint a token
    # that is then offered to a different instance.
    url=$'\e[4;36m'"https://$hostname/-/user_settings/personal_access_tokens"$'\e[0m'
    cat <<EOF
Create or renew a fine-grained PAT at:
  $url

Recommended scopes for a sandboxed Claude Code:
  - api, read_repository, write_repository
  - Short expiration so a leaked token expires quickly

EOF
    read -sp "GitLab PAT for $hostname: " t && echo
    echo "$t" | glab auth login --stdin --hostname "$hostname"
    unset t
    # glab has no `auth login --git-protocol` flag (it does not exist in glab
    # 1.36); passing one aborts the login. git_protocol is config, and its
    # shipped default is ssh — which a sandbox has no key for. Set it per-host
    # so only the instance just authenticated is pinned to https.
    glab config set -h "$hostname" git_protocol https
    # glab auth login records the token against the host but leaves the
    # default host alone — it ships as gitlab.com. Outside a git repo, and in
    # any repo whose remote glab cannot map, bare `glab` commands then query
    # gitlab.com and fail as unauthenticated. Point the default at the
    # instance just authenticated.
    # --global is required: without it `glab config set` writes the
    # repository-local .git/glab-cli/config.yml (leaving the real default at
    # gitlab.com), and outside a git repo it fails outright with "not a git
    # repository" — which under `set -e` would abort us right after login.
    glab config set --global host "$hostname"
    echo "Default glab host set to $hostname."
    glab auth status
}

cmd_update() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "claude-sandbox: update must run as root (install requires it)." >&2
        exit 1
    fi
    if [ "${IS_SANDBOX:-0}" = "1" ]; then
        echo "claude-sandbox: refusing to update from inside a sandboxed claude session." >&2
        exit 1
    fi
    if [ -d /opt/claude-sandbox ]; then
        echo "claude-sandbox: this is the published container image — update by pulling a newer image and recreating the container (uvx claude-sandbox@latest --recreate)." >&2
        exit 1
    fi
    if [ "$(cat "$INSTALLER_FILE" 2>/dev/null)" = uvx ]; then
        # The wheel is the pin (ADR 23): a git clone here would step past it.
        echo "claude-sandbox: this sandbox was installed from the PyPI wheel — update with: uvx claude-sandbox@latest install" >&2
        echo "  (or bump the pinned version in your devcontainer's postCreate)" >&2
        exit 1
    fi
    local tmp
    tmp="$(mktemp -d)"
    git clone --quiet "$REPO_URL" "$tmp/claude-sandbox"
    # Release selection lives in `install` (which defaults to the newest
    # stable tag) so there is ONE implementation of "what is current" —
    # shared with a first install from the documented clone one-liner.
    # This clone is pristine and on the default branch, so the default
    # applies and no flag is needed.
    # Safe self-replacement: install overwrites the very file bash is
    # reading incrementally. exec into a fresh `bash -c` — the command
    # string lives in argv, not in this file — so nothing reads this
    # script again after the overwrite; the temp clone is removed after.
    exec bash -c 'bash "$1/claude-sandbox/install" && rm -rf "$1"' claude-sandbox-update "$tmp"
}

cmd_verify() {
    local agent=claude
    while [ "$#" -gt 0 ]; do
        case "$1" in
            --agent)
                case "${2:-}" in
                    claude|codex|pi) agent="$2"; shift 2 ;;
                    *) echo "claude-sandbox: verify --agent needs claude, codex, or pi" >&2; exit 2 ;;
                esac
                ;;
            *) echo "claude-sandbox: unknown verify option '$1'" >&2; exit 2 ;;
        esac
    done

    if [ "${IS_SANDBOX:-0}" = 1 ]; then
        exec /bin/bash /usr/libexec/claude-sandbox/verify-sandbox-battery.sh
    fi
    exec "/usr/local/bin/$agent" --sandbox-verify
}

# Discover lllm2's single loaded model and actual per-slot context. The manual
# form remains available for servers without llama.cpp's /props endpoint.
cmd_pi_local() (
    local model="" context="" port="${CLAUDE_SANDBOX_LOCAL_MODEL_PORT:-1920}" discover=false dir file tmp
    case "$#" in
        0) discover=true ;;
        2) if [ "$1" = --port ]; then port="$2"; discover=true; else model="$1"; context="$2"; fi ;;
        3) model="$1"; context="$2"; port="$3" ;;
        *) echo 'Usage: claude-sandbox pi-local [--port PORT] or pi-local MODEL CONTEXT [PORT]' >&2; exit 2 ;;
    esac
    if ! [[ "$port" =~ ^[1-9][0-9]{0,4}$ ]] || (( port > 65535 )); then
        echo 'claude-sandbox: local model port must be 1–65535.' >&2
        exit 2
    fi
    if "$discover"; then
        if ! model="$(curl --noproxy '*' -fs --connect-timeout 1 --max-time 3 "http://127.0.0.1:$port/v1/models" \
            | jq -er '.data | select(length == 1) | .[0].id | select(type == "string" and length > 0)')" \
            || ! context="$(curl --noproxy '*' -fs --connect-timeout 1 --max-time 3 "http://127.0.0.1:$port/props" \
            | jq -er '.default_generation_settings.n_ctx | select(type == "number" and . == floor)')"; then
            echo "claude-sandbox: could not discover one loaded model and its context on port $port; existing Pi configuration kept. Start a model in lllm2, or use pi-local MODEL CONTEXT [PORT]." >&2
            exit 1
        fi
    fi
    if [ -z "$model" ] || ! [[ "$context" =~ ^[1-9][0-9]{2,6}$ ]] || (( context < 512 )); then
        echo 'claude-sandbox: supply a model ID and context of 512–9999999 tokens.' >&2
        exit 2
    fi
    umask 077
    dir="$HOME/.pi/agent"
    file="$dir/models.json"
    mkdir -p "$dir"
    tmp="$(mktemp "$dir/.models.XXXXXX")"
    trap 'rm -f "$tmp"' EXIT
    if [ -e "$file" ]; then
        jq -e 'type == "object" and ((.providers // {}) | type == "object")' "$file" >/dev/null
    else
        printf '{}\n' > "$tmp"
        file="$tmp"
    fi
    local result
    result="$(jq --arg model "$model" --argjson context "$context" --arg port "$port" '
        .providers.lllm2 = ((.providers.lllm2 // {}) + {
            baseUrl: ("http://127.0.0.1:" + $port + "/v1"),
            api: "openai-completions", apiKey: (.providers.lllm2.apiKey // "local"),
            compat: ({supportsDeveloperRole: false, supportsReasoningEffort: false} + (.providers.lllm2.compat // {})),
            models: [((.providers.lllm2.models // [] | map(select(.id == $model)) | first) // {}) +
                     {id: $model, name: ("Local (lllm2): " + $model),
                      contextWindow: $context, maxTokens: ([$context / 4 | floor, 32000] | min)}]
        })' "$file")"
    printf '%s\n' "$result" > "$tmp"
    mv "$tmp" "$dir/models.json"
    echo "Configured Pi's lllm2 provider at http://127.0.0.1:$port/v1."
    echo "Select lllm2 in Pi's /model picker, or launch pi --provider lllm2."
    echo "The relay setting in /etc/claude-sandbox.conf must match port $port (shipped default: 1920)."
)

# --- doctor: check, and with --fix apply, the recommended user setup ---------
# Seams for tests only: the paths doctor reads and writes.
LIBEXEC_DIR="${CLAUDE_SANDBOX_LIBEXEC:-/usr/libexec/claude-sandbox}"
TAG_FILE="${CLAUDE_SANDBOX_TAG_FILE:-/etc/claude-sandbox-tag}"
TERMINAL_CONFIG="${USER_TERMINAL_CONFIG:-/user-terminal-config}"
SL_CMD='bash $HOME/.claude/statusline-command.sh'
PROMPT_BEGIN='# >>> claude-sandbox prompt tag >>>'
PROMPT_END='# <<< claude-sandbox prompt tag <<<'

DOCTOR_FIX=0
DOCTOR_PENDING=0
# report STATUS SUBJECT DETAIL: one aligned result line.
report() { printf '  %-8s %-22s %s\n' "$1" "$2" "$3"; }
# pending SUBJECT DETAIL: a problem that --fix would repair.
pending() { report todo "$1" "$2"; DOCTOR_PENDING=1; }

# backup FILE: copy FILE aside with a timestamp and say where it went.
backup() {
    local copy
    copy="$1.bak-$(date +%Y%m%d-%H%M%S)"
    cp -p "$1" "$copy"
    report backup "$(basename "$1")" "saved the original as $copy"
}

# The prompt block for SHELL (zsh or bash). It prefixes the tag in white to the
# prompt inside a launcher-made container and does nothing anywhere else, so
# a terminal-config shared with the host or a devcontainer stays safe. The
# case guard stops a re-sourced rc file adding the prefix twice.
prompt_block() {
    printf '%s\n' "$PROMPT_BEGIN"
    printf '%s\n' '# Added by `claude-sandbox doctor --fix`: shows the tag of the claude-sandbox' \
        '# container this shell runs in. Delete this block to remove it.'
    case "$1" in
        zsh) cat <<'EOF'
if [ -r /etc/claude-sandbox-tag ]; then
    __cs_tag="%F{white}$(cat /etc/claude-sandbox-tag)%f "
    # Put the tag at the start of the line above the input line. A theme such
    # as dst opens with a blank line and puts user@host on the next one; the
    # tag belongs on that user@host line, not the blank one.
    case "$PROMPT" in
        *"$__cs_tag"*) ;;
        *$'\n'*$'\n'*)
            __cs_head="${PROMPT%$'\n'*$'\n'*}"
            PROMPT="$__cs_head"$'\n'"$__cs_tag${PROMPT#"$__cs_head"$'\n'}"
            unset __cs_head ;;
        *) PROMPT="$__cs_tag$PROMPT" ;;
    esac
    unset __cs_tag
fi
EOF
        ;;
        bash) cat <<'EOF'
if [ -r /etc/claude-sandbox-tag ]; then
    __cs_tag="$(cat /etc/claude-sandbox-tag)"
    case "$PS1" in *"$__cs_tag"*) ;; *) PS1="\[\033[0;37m\]$__cs_tag\[\033[0m\] $PS1" ;; esac
    unset __cs_tag
fi
EOF
        ;;
    esac
    printf '%s\n' "$PROMPT_END"
}

doctor_tag() {
    if [ -s "$TAG_FILE" ]; then
        report ok "container tag" "$(cat "$TAG_FILE")"
    else
        # Not fixable from inside: the launcher sets the tag at create time.
        report info "container tag" "none; containers from the uvx launcher get one (use --recreate on an older one)"
    fi
}

# doctor_file SUBJECT SHIPPED DEST [MODE]: DEST must be a copy of SHIPPED.
doctor_file() {
    local subject="$1" shipped="$2" dest="$3" mode="${4:-0755}"
    if [ ! -r "$shipped" ]; then
        report skip "$subject" "$shipped is missing; re-run the install"
    elif [ -f "$dest" ] && cmp -s "$shipped" "$dest"; then
        report ok "$subject" "$dest"
    elif [ "$DOCTOR_FIX" = 0 ]; then
        pending "$subject" "$dest is $([ -f "$dest" ] && echo 'not the recommended one' || echo absent)"
    else
        mkdir -p "$(dirname "$dest")"
        [ -f "$dest" ] && backup "$dest"
        install -m "$mode" "$shipped" "$dest"
        report fixed "$subject" "installed $dest"
    fi
}

doctor_claude_settings() {
    local settings="$HOME/.claude/settings.json" input tmp
    if [ -f "$settings" ] && ! jq -e 'type == "object"' "$settings" >/dev/null 2>&1; then
        report skip "claude settings" "$settings is not a JSON object; set statusLine by hand"
    elif [ -f "$settings" ] && jq -e --arg c "$SL_CMD" '.statusLine.command == $c' "$settings" >/dev/null; then
        report ok "claude settings" "statusLine runs the script"
    elif [ "$DOCTOR_FIX" = 0 ]; then
        pending "claude settings" "statusLine in $settings does not run the script"
    else
        input='{}'
        if [ -f "$settings" ]; then
            backup "$settings"
            input="$(cat "$settings")"
        fi
        mkdir -p "$(dirname "$settings")"
        tmp="$(mktemp "$settings.XXXXXX")"
        jq --arg c "$SL_CMD" '.statusLine = {type: "command", command: $c}' <<<"$input" > "$tmp"
        chmod 0644 "$tmp"
        mv "$tmp" "$settings"
        report fixed "claude settings" "statusLine now runs the script"
    fi
}

# prompt_block_in RC: the tag block already in RC, markers included.
prompt_block_in() {
    awk -v b="$PROMPT_BEGIN" -v e="$PROMPT_END" '$0 == b { p = 1 } p { print } $0 == e { p = 0 }' "$1"
}

# doctor_prompt SHELL: the prompt block in the shared rc file for SHELL.
doctor_prompt() {
    local rc="$TERMINAL_CONFIG/${1}rc"
    if [ ! -f "$rc" ]; then
        report skip "$1 prompt" "$rc is absent (only DLS base-image shells read it)"
    elif grep -qxF "$PROMPT_BEGIN" "$rc"; then
        if [ "$(prompt_block_in "$rc")" = "$(prompt_block "$1")" ]; then
            report ok "$1 prompt" "$rc"
        elif [ "$DOCTOR_FIX" = 0 ]; then
            pending "$1 prompt" "$rc has an older tag block"
        else
            local block
            backup "$rc"
            block="$(mktemp)"
            prompt_block "$1" > "$block"
            # cat into the rc file, not mv: keep its mode and any symlink.
            awk -v b="$PROMPT_BEGIN" -v e="$PROMPT_END" -v f="$block" '
                $0 == b { while ((getline line < f) > 0) print line; skip = 1; next }
                skip && $0 == e { skip = 0; next }
                !skip { print }' "$rc" > "$block.rc"
            cat "$block.rc" > "$rc"
            rm -f "$block" "$block.rc"
            report fixed "$1 prompt" "updated the tag block in $rc (new shells show it)"
        fi
    elif [ "$DOCTOR_FIX" = 0 ]; then
        pending "$1 prompt" "$rc has no container tag in the prompt"
    else
        backup "$rc"
        { printf '\n'; prompt_block "$1"; } >> "$rc"
        report fixed "$1 prompt" "appended the tag block to $rc (new shells show it)"
    fi
}

cmd_doctor() {
    case "${1:-}" in
        '') ;;
        --fix) DOCTOR_FIX=1 ;;
        *) echo "claude-sandbox: doctor takes only --fix" >&2; exit 2 ;;
    esac
    if [ "$DOCTOR_FIX" = 1 ] && [ "${IS_SANDBOX:-0}" = 1 ]; then
        echo "claude-sandbox: refusing doctor --fix inside a sandboxed agent session; run it from a shell." >&2
        exit 1
    fi
    doctor_tag
    doctor_file "claude status line" "$LIBEXEC_DIR/statusline-command.sh" "$HOME/.claude/statusline-command.sh"
    doctor_claude_settings
    doctor_file "pi footer" "$LIBEXEC_DIR/pi-sandbox-tag.ts" "$HOME/.pi/agent/extensions/claude-sandbox-tag.ts" 0644
    doctor_prompt zsh
    doctor_prompt bash
    if [ "$DOCTOR_PENDING" = 1 ]; then
        echo "Run \`claude-sandbox doctor --fix\` to apply the recommended setup. It backs up every file it changes."
        exit 1
    fi
}

cmd_version() {
    if [ -r "$VERSION_FILE" ]; then
        echo "claude-sandbox $(cat "$VERSION_FILE")"
    else
        echo "claude-sandbox: version unknown ($VERSION_FILE missing — re-run install)" >&2
        exit 1
    fi
}

case "${1:-help}" in
    install)
        echo "claude-sandbox is already installed in this container. No installation is needed."
        echo "Run claude, codex or pi to start a sandboxed agent session."
        ;;
    gh-auth)              shift; cmd_gh_auth "$@" ;;
    glab-auth)            shift; cmd_glab_auth "$@" ;;
    update)               shift; cmd_update "$@" ;;
    verify)               shift; cmd_verify "$@" ;;
    pi-local)             shift; cmd_pi_local "$@" ;;
    doctor)               shift; cmd_doctor "$@" ;;
    version|--version|-v) cmd_version ;;
    help|-h|--help)       usage ;;
    *)
        echo "claude-sandbox: unknown command '$1'" >&2
        usage >&2
        exit 2
        ;;
esac
