#!/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.
#
# Usage: ./run.sh [--build] [--check-only] [--workspace DIR]
#   --build          rebuild the image even if it exists
#   --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}"
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
while (( $# )); do
    case "$1" in
        --build) build=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,10p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) echo "Unknown argument: $1" >&2; exit 2 ;;
    esac
    shift
done

die() { echo "error: $*" >&2; exit 1; }
info() { echo "==> $*"; }

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

# 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 ----------------------------------------------------------
if "$RUNTIME" container inspect "$CONTAINER" >/dev/null 2>&1; then
    die "container '$CONTAINER' already exists; stop it with: $RUNTIME stop $CONTAINER"
fi

# --- 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"

# 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.
info "Starting Claude Desktop and Ghidra in container '$CONTAINER' (workspace: $WORKSPACE -> /workspace)"
"$RUNTIME" run -d "${run_args[@]}" \
    --name "$CONTAINER" \
    -v "$WORKSPACE:/workspace" \
    "$IMAGE" >/dev/null
info "Running. Logs: $RUNTIME logs -f $CONTAINER   Stop: $RUNTIME stop $CONTAINER"