#!/usr/bin/env bash
# Start the Claude Desktop container, forwarding the host's X server.
# Before launching, verify that a window from the container can actually be
# drawn on the host display. Asks (in a folder picker) which host folder to
# share with Claude; it is mounted read-write at /workspace in the container.
# Also starts the REA MCP server in a second, isolated container and stops it
# when Claude Desktop's container ends.
#
# Usage: ./run.sh [--build] [--build-only] [--check-only] [--workspace DIR]
#   --build          rebuild the image even if it exists
#   --build-only     (re)build the image and exit
#   --check-only     run the display check and exit without starting Claude
#   --workspace DIR  share DIR without asking (also: WORKSPACE=DIR)

set -euo pipefail

IMAGE="${IMAGE:-claude-re-setup:latest}"
CONTAINER="${CONTAINER:-claude-desktop}"
REA_CONTAINER="${REA_CONTAINER:-claude-re-mcp}"
# Internal network (no route out) shared by Claude Desktop and REA.
NETWORK="${NETWORK:-claude-re-setup-internal}"
RUNTIME="${RUNTIME:-docker}"
WORKSPACE="${WORKSPACE:-}"

script_dir="$(dirname "$(readlink -f "$0")")"
state_dir="${XDG_STATE_HOME:-$HOME/.local/state}/claude-re-setup"
last_workspace_file="$state_dir/last_workspace"

build=0
check_only=0
build_only=0
while (( $# )); do
    case "$1" in
        --build) build=1 ;;
        --build-only) build=1 build_only=1 ;;
        --check-only) check_only=1 ;;
        --workspace) [[ $# -ge 2 ]] || { echo "--workspace needs a directory" >&2; exit 2; }
                     WORKSPACE="$2"; shift ;;
        --workspace=*) WORKSPACE="${1#*=}" ;;
        -h|--help) sed -n '2,13p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) echo "Unknown argument: $1" >&2; exit 2 ;;
    esac
    shift
done

# Started from the app menu (install.sh) there is no terminal to read
# errors in, so they are also shown in a dialog.
die() {
    echo "error: $*" >&2
    if [[ ! -t 2 ]] && command -v zenity >/dev/null; then
        zenity --error --no-wrap --title="Claude RE workstation" --text="$*" 2>/dev/null || true
    fi
    exit 1
}
info() { echo "==> $*"; }

command -v "$RUNTIME" >/dev/null || die "'$RUNTIME' not found (set RUNTIME=podman to use podman)"

# --- Image build -------------------------------------------------------------
ensure_image() {
    if (( build )) || ! "$RUNTIME" image inspect "$IMAGE" >/dev/null 2>&1; then
        info "Building $IMAGE"
        "$RUNTIME" build \
            --build-arg USER_UID="$(id -u)" \
            --build-arg USER_GID="$(id -g)" \
            -t "$IMAGE" "$script_dir"
    fi
}
if (( build_only )); then
    ensure_image
    info "Built $IMAGE"
    exit 0
fi

command -v xauth >/dev/null || die "xauth not found on the host (dnf install xorg-x11-xauth)"

# --- Display -----------------------------------------------------------------
# On Wayland sessions this is XWayland, which GNOME/KDE start on demand.
[[ -n "${DISPLAY:-}" ]] || die "DISPLAY is not set; run this from a graphical session"

display_num="${DISPLAY#*:}"
display_num="${display_num%%.*}"
[[ -S "/tmp/.X11-unix/X${display_num}" ]] \
    || die "X socket /tmp/.X11-unix/X${display_num} not found for DISPLAY=$DISPLAY"

# Hand the container a copy of this display's cookie, rewritten to match any
# hostname (family ffff). This avoids opening the server up with 'xhost +'.
xauth_dir="${XDG_RUNTIME_DIR:-/tmp}/claude-re-setup"
xauth_file="$xauth_dir/Xauthority"
mkdir -p "$xauth_dir"
chmod 700 "$xauth_dir"
: > "$xauth_file"
xauth nlist "$DISPLAY" | sed -e 's/^..../ffff/' | xauth -f "$xauth_file" nmerge - \
    || die "could not export the X cookie for $DISPLAY"
[[ -s "$xauth_file" ]] || die "no X cookie found for $DISPLAY (is XAUTHORITY set?)"
# The container user has the same UID as us, so 0600 is enough.
chmod 600 "$xauth_file"

# --- Image -------------------------------------------------------------------
ensure_image

# label=disable: the X socket and cookie are host files that SELinux would
# otherwise block, and relabeling /tmp/.X11-unix with :z would break the host.
run_args=(
    --rm
    --init
    --security-opt label=disable
    --shm-size 1g
    -e DISPLAY=":${display_num}"
    -e XAUTHORITY=/tmp/.Xauthority
    -v /tmp/.X11-unix:/tmp/.X11-unix:ro
    -v "$xauth_file:/tmp/.Xauthority:ro"
)

# --- Display check -----------------------------------------------------------
# 1. xdpyinfo proves the container can connect and authenticate to the server.
# 2. xmessage opens a real window; xwininfo confirms the server mapped it.
info "Checking that the container can draw on DISPLAY=$DISPLAY"
check_script='
set -e
xdpyinfo >/dev/null 2>&1 || { echo "FAIL: cannot connect to X server at $DISPLAY"; exit 1; }
echo "ok: connected to X server ($(xdpyinfo | sed -n "s/^vendor string: *//p"))"

title="claude-re-setup-x11-check"
xmessage -title "$title" -center -timeout 3 "Display check from the container: OK" &
for _ in $(seq 1 30); do
    if xwininfo -name "$title" 2>/dev/null | grep -q "Map State: IsViewable"; then
        echo "ok: test window is mapped and viewable"
        wait || true
        exit 0
    fi
    sleep 0.1
done
kill %1 2>/dev/null || true
echo "FAIL: test window was created but never became viewable"
exit 1
'
if ! "$RUNTIME" run "${run_args[@]}" "$IMAGE" bash -c "$check_script"; then
    die "display check failed; see the output above"
fi

(( check_only )) && { info "Display check passed"; exit 0; }

# --- Claude Desktop ----------------------------------------------------------
for c in "$CONTAINER" "$REA_CONTAINER"; do
    if "$RUNTIME" container inspect "$c" >/dev/null 2>&1; then
        die "container '$c' already exists; stop it with: $RUNTIME stop $c"
    fi
done

# --- Model provider ----------------------------------------------------------
# Asked on every start. Left empty, Claude Desktop starts normally (claude.ai
# sign-in, Anthropic's models). Given a base URL, Claude Desktop runs in its
# third-party "gateway" mode: Chat, Code and Cowork send their requests to
# that URL, which must speak the Anthropic Messages API (/v1/messages), e.g.
# a self-hosted model behind LiteLLM. The settings go into Claude Desktop's
# managed configuration file (/etc/claude-desktop/managed-settings.json on
# Linux). Claude Desktop ignores that file unless root owns it and only root
# can write it, so it is copied into the created container (as root, 0644)
# just before it starts, rather than mounted; the copy on this side, in a
# private runtime folder, is deleted right after. Nothing is remembered
# between starts, the token included.
managed_settings_dir="$xauth_dir/managed-settings"
rm -rf "$managed_settings_dir"
gateway_url=""
gateway_token=""
gateway_scheme="bearer"
gateway_models=""

# JSON string literal for $1 (refuses control characters).
json_str() {
    [[ "$1" != *[[:cntrl:]]* ]] || die "control characters are not allowed in: $1"
    local s="${1//\\/\\\\}"
    printf '"%s"' "${s//\"/\\\"}"
}

provider_text="Leave the base URL empty to use Anthropic's models (sign in to claude.ai as usual).

To use another model, enter the base URL of a server that speaks the Anthropic Messages API, without /v1 (e.g. http://host.docker.internal:4000 for a LiteLLM proxy on this machine). With http:// the token is sent unencrypted. Models: comma-separated IDs, the first is the default; leave empty to ask the server's /v1/models."
if command -v zenity >/dev/null; then
    info "Choose the model provider in the dialog"
    form="$(zenity --forms --title="Claude RE workstation: model provider" \
        --text="$provider_text" --separator=$'\x1f' \
        --add-entry="Base URL (ANTHROPIC_BASE_URL)" \
        --add-password="Token" \
        --add-combo="Send token as" --combo-values="Authorization: Bearer|x-api-key" \
        --add-entry="Models (optional)" 2>/dev/null)" \
        || die "model provider dialog cancelled; not starting Claude"
    IFS=$'\x1f' read -r gateway_url gateway_token scheme_choice gateway_models <<< "$form"
    [[ "${scheme_choice:-}" == "x-api-key" ]] && gateway_scheme="x-api-key"
elif [[ -t 0 ]]; then
    echo "$provider_text"
    read -r -p "Base URL (empty = Anthropic): " gateway_url
    if [[ -n "$gateway_url" ]]; then
        read -r -s -p "Token: " gateway_token; echo
        read -r -p "Models (optional): " gateway_models
    fi
fi

# Trim spaces and a trailing slash or /v1 (Claude Desktop appends /v1/messages).
gateway_url="$(sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' -e 's#/*$##' -e 's#/v1$##' <<< "$gateway_url")"
gateway_token="$(sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' <<< "$gateway_token")"
if [[ -n "$gateway_url" ]]; then
    [[ "$gateway_url" =~ ^https?://[^/[:space:]]+ ]] \
        || die "the base URL must start with http:// or https://: $gateway_url"
    # localhost in the container is the container itself; a server on this
    # machine is reached through host.docker.internal (added to the container
    # below). It has to listen on an address Docker's bridge can reach, not
    # only 127.0.0.1.
    if [[ "$gateway_url" =~ ^(https?://)(localhost|127\.0\.0\.1|\[::1\])(:|/|$)(.*)$ ]]; then
        gateway_url="${BASH_REMATCH[1]}host.docker.internal${BASH_REMATCH[3]}${BASH_REMATCH[4]}"
        info "Using $gateway_url (localhost inside the container is the container itself)"
    fi
    gateway_shown="$gateway_url"
    # Claude Desktop accepts plain http only for loopback addresses. For an
    # http:// server elsewhere, socat in the container listens on
    # 127.0.0.1:18080 and forwards to it (started by launch-apps.sh), and
    # Claude Desktop is pointed there. The token still crosses the network
    # unencrypted, as it would with any plain-HTTP server.
    if [[ "$gateway_url" =~ ^http://([^/:]+)(:([0-9]+))?(/.*)?$ ]]; then
        fwd_target="${BASH_REMATCH[1]}:${BASH_REMATCH[3]:-80}"
        run_args+=(-e "CLAUDE_RE_GATEWAY_FORWARD=$fwd_target")
        gateway_url="http://127.0.0.1:18080${BASH_REMATCH[4]}"
        info "Plain HTTP: forwarding 127.0.0.1:18080 in the container to $fwd_target (token not encrypted)"
    elif [[ "$gateway_url" == http://* ]]; then
        die "unsupported http:// address (use a hostname or IPv4 address, or https://): $gateway_url"
    fi
    # The key is required in gateway mode; local servers usually ignore it.
    [[ -n "$gateway_token" ]] || gateway_token="none"

    models_json=""
    IFS=',' read -r -a model_list <<< "$gateway_models"
    for m in "${model_list[@]}"; do
        m="$(sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' <<< "$m")"
        [[ -n "$m" ]] && models_json+="${models_json:+, }$(json_str "$m")"
    done

    # Removed again if the start fails before it is copied in.
    trap 'rm -rf "$managed_settings_dir"' EXIT
    mkdir -m 700 "$managed_settings_dir"
    ( umask 077
      {
        echo "{"
        echo "  \"inferenceProvider\": \"gateway\","
        echo "  \"inferenceGatewayBaseUrl\": $(json_str "$gateway_url"),"
        echo "  \"inferenceGatewayApiKey\": $(json_str "$gateway_token"),"
        echo "  \"inferenceGatewayAuthScheme\": \"$gateway_scheme\","
        [[ -n "$models_json" ]] && echo "  \"inferenceModels\": [$models_json],"
        echo "  \"disableDeploymentModeChooser\": true"
        echo "}"
      } > "$managed_settings_dir/managed-settings.json" )
    info "Model provider: gateway at $gateway_shown"
else
    info "Model provider: Anthropic (claude.ai sign-in)"
fi
unset gateway_token
# For self-hosted servers on this machine (see above).
run_args+=(--add-host host.docker.internal:host-gateway)

# --- Workspace ---------------------------------------------------------------
# The folder Claude can read and write. Chosen in the desktop's own folder
# picker (zenity on GNOME, kdialog on KDE), else asked in the terminal.
# Default: the last one used, else ./workspace next to this script.
ws_title="Choose the folder Claude can access (mounted at /workspace)"

# Prints the reason a folder can't be shared, or nothing if it's fine.
workspace_problem() {
    case "$1" in
        /|"$HOME") echo "Sharing $1 would give Claude far more than intended; pick a dedicated folder." ;;
        *) [[ -w "$1" ]] || echo "$1 is not writable by $(id -un)." ;;
    esac
}

# Shows a folder picker starting in $1; prints the choice, fails on cancel.
gui_pick() {
    if command -v zenity >/dev/null; then
        zenity --file-selection --directory --title="$ws_title" --filename="$1/" 2>/dev/null
    elif command -v kdialog >/dev/null; then
        kdialog --title "$ws_title" --getexistingdirectory "$1" 2>/dev/null
    else
        return 2
    fi
}
gui_error() {
    if command -v zenity >/dev/null; then zenity --error --no-wrap --text="$1" 2>/dev/null
    else kdialog --error "$1" 2>/dev/null; fi
}

if [[ -z "$WORKSPACE" ]]; then
    default_ws="$script_dir/workspace"
    [[ -r "$last_workspace_file" ]] && default_ws="$(<"$last_workspace_file")"
    # The picker needs an existing folder to open in. ./workspace is
    # git-ignored, so creating it is harmless; a stale last choice falls back.
    [[ -d "$default_ws" ]] || { default_ws="$script_dir/workspace"; mkdir -p "$default_ws"; }

    if command -v zenity >/dev/null || command -v kdialog >/dev/null; then
        info "Choose the folder to share in the dialog"
        while :; do
            WORKSPACE="$(gui_pick "$default_ws")" || die "no folder chosen; not starting Claude"
            problem="$(workspace_problem "$(readlink -f "$WORKSPACE")")"
            [[ -z "$problem" ]] && break
            gui_error "$problem"
        done
    elif [[ -t 0 ]]; then
        echo "Claude will get read-write access to one folder, mounted at /workspace."
        read -r -e -p "Folder to share [$default_ws]: " WORKSPACE
    fi
    WORKSPACE="${WORKSPACE:-$default_ws}"
fi
# Expand a leading ~ the way the shell would have.
[[ "$WORKSPACE" == "~" || "$WORKSPACE" == "~/"* ]] && WORKSPACE="$HOME${WORKSPACE#\~}"

if [[ ! -d "$WORKSPACE" ]]; then
    [[ -e "$WORKSPACE" ]] && die "$WORKSPACE exists but is not a directory"
    if [[ -t 0 ]]; then
        read -r -p "$WORKSPACE does not exist. Create it? [Y/n] " answer
        [[ "${answer:-y}" =~ ^[Yy] ]] || die "no folder to share"
    fi
    mkdir -p "$WORKSPACE" || die "could not create $WORKSPACE"
fi
WORKSPACE="$(readlink -f "$WORKSPACE")"

problem="$(workspace_problem "$WORKSPACE")"
[[ -z "$problem" ]] || die "$problem"

mkdir -p "$state_dir"
printf '%s\n' "$WORKSPACE" > "$last_workspace_file"

# --- REA MCP server ----------------------------------------------------------
# REA (Reverse Engineer Anything) gets its own container from the same image,
# because some of its tools run the target: no display, no capabilities, only
# the workspace mounted, and only the internal network, which has no route
# out. It serves MCP over HTTP as "rea:8765"; Claude Desktop reaches it
# through mcp-remote (see claude-desktop/claude_desktop_config.json).
if ! "$RUNTIME" network inspect "$NETWORK" >/dev/null 2>&1; then
    info "Creating internal network $NETWORK"
    "$RUNTIME" network create --internal "$NETWORK" >/dev/null
fi

ws_mount=(-v "$WORKSPACE:/workspace")
info "Starting the REA MCP server in container '$REA_CONTAINER'"
"$RUNTIME" run -d --rm --init \
    --name "$REA_CONTAINER" \
    --network "$NETWORK" --network-alias rea \
    --cap-drop ALL --security-opt no-new-privileges \
    --security-opt label=disable \
    "${ws_mount[@]}" \
    "$IMAGE" rea mcp --http --host 0.0.0.0 --port 8765 >/dev/null
stop_rea() { "$RUNTIME" stop "$REA_CONTAINER" >/dev/null 2>&1 || true; }

# --- Claude Desktop ----------------------------------------------------------
# No config volume: the container is removed on stop (--rm), so the sign-in
# is discarded and Claude asks for a login on every start.
# No :z on the workspace mount: label=disable already lets the container use
# it, and relabeling an arbitrary host folder could break other apps.
# Created first, so it is on the internal network (as well as the default one,
# for sign-in) before it starts.
info "Starting Claude Desktop and Ghidra in container '$CONTAINER' (workspace: $WORKSPACE -> /workspace)"
# Copies the provider settings (if any) into the created container, owned
# by root and readable by Claude Desktop, then deletes them on this side.
copy_managed_settings() {
    [[ -d "$managed_settings_dir" ]] || return 0
    tar -C "$managed_settings_dir" --owner=0 --group=0 --numeric-owner --mode=0644 \
        -cf - managed-settings.json \
        | "$RUNTIME" cp - "$CONTAINER:/etc/claude-desktop/" \
        && rm -rf "$managed_settings_dir"
}
if ! { "$RUNTIME" create "${run_args[@]}" --name "$CONTAINER" "${ws_mount[@]}" "$IMAGE" >/dev/null \
       && "$RUNTIME" network connect "$NETWORK" "$CONTAINER" \
       && copy_managed_settings \
       && "$RUNTIME" start "$CONTAINER" >/dev/null; }; then
    "$RUNTIME" rm -f "$CONTAINER" >/dev/null 2>&1 || true
    stop_rea
    die "could not start Claude Desktop"
fi
trap - EXIT

# Stop REA once Claude Desktop's container is gone, in the background and
# detached from this terminal.
setsid nohup bash -c '"$1" wait "$2" >/dev/null 2>&1; "$1" stop "$3" >/dev/null 2>&1' \
    _ "$RUNTIME" "$CONTAINER" "$REA_CONTAINER" </dev/null >/dev/null 2>&1 &

info "Running. Logs: $RUNTIME logs -f $CONTAINER (REA: $REA_CONTAINER)   Stop: $RUNTIME stop $CONTAINER"