#!/bin/bash
# Author: Koushik Sen (ksen@berkeley.edu)
# Contributors:
# Koushik Sen (ksen@berkeley.edu)
# add your name here
# Run KISS Sorcar on a remote Linux machine over SSH and use it from your
# local VS Code (browser-based code-server).
#
# Usage:  ./sorcar-linux [--force] user@ip-address [PORT]
#
# Flow:
#   1. SSH into user@ip-address.
#   1b. Idempotency probe: if the host already has code-server, ~/kiss_ai and
#       the KISS Sorcar extension, skip reinstalling — either reuse the
#       already-running code-server or just relaunch it.  Pass --force to
#       reprovision anyway.
#   2. Install code-server (https://code-server.dev) if it is missing.
#   3. Install KISS Sorcar (kiss_ai) with the README curl one-liner
#      ( curl -fsSL .../scripts/install.sh | bash ) — only if ~/kiss_ai is
#      not already present.  A tiny `code` shim makes the installer drop the
#      KISS Sorcar extension straight into code-server instead of a desktop
#      VS Code (which a headless server does not have).
#   4. Launch code-server on the remote host bound to 127.0.0.1:PORT.
#   5. Open an SSH tunnel localhost:PORT -> remote 127.0.0.1:PORT and point
#      your local browser at it.  code-server *is* VS Code running in the
#      browser, so this is your local VS Code talking to the remote machine.
#
# The KISS Sorcar runtime (uv, Python deps, Playwright, API keys, etc.) is
# finished by the extension itself once code-server opens, exactly like the
# desktop and Docker installs.
set -euo pipefail

# ---------------------------------------------------------------------------
# Logging helpers (mirrors sorcar-docker)
# ---------------------------------------------------------------------------
info() { printf '\033[0;32m[INFO]\033[0m  %s\n' "$*"; }
step() { printf '\033[0;34m[STEP]\033[0m  %s\n' "$*"; }
warn() { printf '\033[1;33m[WARN]\033[0m  %s\n' "$*"; }
die()  { printf '\033[0;31m[ERR]\033[0m  %s\n' "$*" >&2; exit 1; }

# ---------------------------------------------------------------------------
# 1. Parse arguments
# ---------------------------------------------------------------------------
FORCE=0
ARGS=()
for arg in "$@"; do
    case "$arg" in
        --force|-f) FORCE=1 ;;
        *) ARGS+=("$arg") ;;
    esac
done
TARGET="${ARGS[0]:-}"
PORT="${ARGS[1]:-8080}"

if [[ -z "$TARGET" ]]; then
    die "Usage: $0 [--force] user@ip-address [PORT]"
fi
if [[ "$TARGET" != *"@"* ]]; then
    warn "TARGET '$TARGET' has no 'user@' part; using the SSH default user."
fi
URL="http://localhost:${PORT}"

# ---------------------------------------------------------------------------
# 2. Pre-flight checks
# ---------------------------------------------------------------------------
command -v ssh &>/dev/null  || die "ssh is not installed. Install OpenSSH first."
command -v curl &>/dev/null || die "curl is not installed. Install curl first."

step "Checking SSH connectivity to $TARGET..."
ssh -o ConnectTimeout=15 -o BatchMode=no "$TARGET" 'echo ok' >/dev/null \
    || die "Cannot SSH into $TARGET. Check the address and your SSH keys."
info "SSH connection OK"

# ---------------------------------------------------------------------------
# 2b. Idempotency probe: is the host already provisioned?
#
# Classifies the remote host as one of:
#   provisioned-running : code-server + ~/kiss_ai + KISS Sorcar extension are
#                         all present AND code-server is already answering on
#                         127.0.0.1:$PORT -> skip provisioning entirely.
#   provisioned-stopped : everything is installed but code-server is not
#                         serving on $PORT -> skip installs, only relaunch.
#   unprovisioned       : something is missing -> run the full setup.
# --force skips the probe and always runs the full setup.
# ---------------------------------------------------------------------------
if [[ "$FORCE" == "1" ]]; then
    STATUS="forced"
    warn "--force given: reprovisioning $TARGET even if already set up."
else
    step "Checking whether $TARGET is already provisioned..."
    STATUS=$(ssh "$TARGET" "PORT='$PORT' bash -s" <<'PROBE' || echo unprovisioned
PORT="${PORT:-8080}"
export PATH="$HOME/.local/bin:$PATH"
command -v code-server >/dev/null 2>&1        || { echo unprovisioned; exit 0; }
[ -d "$HOME/kiss_ai" ]                        || { echo unprovisioned; exit 0; }
code-server --list-extensions 2>/dev/null | grep -qi 'ksenxx.kiss-sorcar' \
                                              || { echo unprovisioned; exit 0; }
if curl -sf "http://127.0.0.1:$PORT/healthz" -o /dev/null 2>/dev/null \
   || curl -sf "http://127.0.0.1:$PORT" -o /dev/null 2>/dev/null; then
    echo provisioned-running
else
    echo provisioned-stopped
fi
PROBE
    )
fi

# ---------------------------------------------------------------------------
# 3. Remote bootstrap: code-server + kiss_ai + launch code-server.
#    Skipped entirely when the probe found a healthy, already-running
#    install; runs in relaunch-only mode (SKIP_INSTALL=1) when everything is
#    installed but code-server is down.
#
# The heredoc is single-quoted so nothing is expanded locally; PORT and
# SKIP_INSTALL are passed through the environment on the ssh command line.
# ---------------------------------------------------------------------------
if [[ "$STATUS" == "provisioned-running" ]]; then
    info "$TARGET is already provisioned: code-server + kiss_ai + extension found,"
    info "code-server already serving on 127.0.0.1:$PORT — skipping reinstall."
    info "(Use --force to reprovision from scratch.)"
else
    SKIP_INSTALL=0
    if [[ "$STATUS" == "provisioned-stopped" ]]; then
        SKIP_INSTALL=1
        info "$TARGET already has code-server + kiss_ai — skipping reinstall,"
        info "only relaunching code-server on 127.0.0.1:$PORT."
        step "Restarting code-server on $TARGET..."
    else
        step "Provisioning KISS Sorcar on $TARGET (this can take several minutes)..."
    fi
    ssh "$TARGET" "PORT='$PORT' SKIP_INSTALL='$SKIP_INSTALL' bash -s" <<'REMOTE'
set -e
PORT="${PORT:-8080}"
SKIP_INSTALL="${SKIP_INSTALL:-0}"
WORKSPACE="$HOME/kiss_ai"
export PATH="$HOME/.local/bin:$PATH"
LOG_DIR="$HOME/.kiss"
mkdir -p "$LOG_DIR" "$HOME/.local/bin"

rinfo() { printf '\033[0;36m[%s]\033[0m %s\n' "$(hostname)" "$*"; }

if [ "$SKIP_INSTALL" = "1" ]; then
    rinfo "Already provisioned — skipping code-server/kiss_ai install steps."
else

# --- 3a. code-server -------------------------------------------------------
# Use the standalone method (installs into ~/.local, no sudo) because this
# non-interactive SSH session has no TTY for a sudo password prompt.
if ! command -v code-server &>/dev/null; then
    rinfo "Installing code-server (standalone, no sudo)..."
    curl -fsSL https://code-server.dev/install.sh \
        | sh -s -- --method=standalone --prefix="$HOME/.local"
else
    rinfo "code-server already installed."
fi
command -v code-server &>/dev/null || { echo "ERROR: code-server install failed" >&2; exit 1; }

# --- 3b. `code` shim -------------------------------------------------------
# The README installer (install.sh) looks for a `code` CLI to install the
# extension and to launch the editor.  On a headless server there is no
# desktop VS Code, so we provide a shim:
#   * extension/version subcommands are forwarded to code-server, so the
#     installer drops the KISS Sorcar extension into code-server;
#   * a plain `code <dir>` launch request is swallowed (exit 0) because we
#     start code-server ourselves below with the flags we want.
cat > "$HOME/.local/bin/code" <<'SHIM'
#!/bin/bash
case "${1:-}" in
    --install-extension|--uninstall-extension|--list-extensions|--version|--locate-extension)
        exec code-server "$@" ;;
    *)
        exit 0 ;;
esac
SHIM
chmod +x "$HOME/.local/bin/code"

# --- 3c. Install kiss_ai (README curl install) only if missing -------------
if [ ! -d "$HOME/kiss_ai" ]; then
    rinfo "kiss_ai not found — installing via README curl one-liner..."
    curl -fsSL https://raw.githubusercontent.com/ksenxx/kiss_ai/main/scripts/install.sh | bash
else
    rinfo "kiss_ai already installed at ~/kiss_ai — skipping install."
fi
[ -d "$HOME/kiss_ai" ] || { echo "ERROR: kiss_ai install failed (~/kiss_ai missing)" >&2; exit 1; }

# --- 3d. code-server settings: never block on workspace-trust prompts ------
SETTINGS_DIR="$HOME/.local/share/code-server/User"
mkdir -p "$SETTINGS_DIR"
if [ ! -f "$SETTINGS_DIR/settings.json" ]; then
    cat > "$SETTINGS_DIR/settings.json" <<'JSON'
{
    "security.workspace.trust.enabled": false,
    "security.workspace.trust.startupPrompt": "never",
    "security.workspace.trust.banner": "never",
    "security.workspace.trust.emptyWindow": false
}
JSON
fi

# --- 3e. Make sure the extension is in code-server -------------------------
# install.sh installs it via the shim, but on a re-run (kiss_ai already
# present) the VSIX may have been removed, so guard with a presence check.
if ! code-server --list-extensions 2>/dev/null | grep -qi 'ksenxx.kiss-sorcar'; then
    VSIX="$WORKSPACE/src/kiss/agents/vscode/kiss-sorcar.vsix"
    if [ -f "$VSIX" ]; then
        rinfo "Installing KISS Sorcar extension into code-server..."
        code-server --install-extension "$VSIX" --force || true
    else
        rinfo "WARNING: KISS Sorcar extension not found in code-server and no VSIX present."
    fi
fi

fi # end of SKIP_INSTALL guard around install steps 3a-3e

# --- 3f. (Re)launch code-server on 127.0.0.1:$PORT -------------------------
# Kill any code-server we previously started on this port, then start a fresh
# detached instance that survives the SSH session closing.
pkill -f "code-server.*--bind-addr 127.0.0.1:$PORT" 2>/dev/null || true
sleep 1
rinfo "Launching code-server on 127.0.0.1:$PORT (workspace: $WORKSPACE)..."
setsid bash -c "exec code-server \
    --bind-addr 127.0.0.1:$PORT \
    --auth none \
    --disable-workspace-trust \
    --enable-proposed-api ksenxx.kiss-sorcar \
    '$WORKSPACE'" > "$LOG_DIR/code-server.log" 2>&1 < /dev/null &

# Wait for it to start answering locally before we hand back to the client.
for i in $(seq 1 60); do
    if curl -sf "http://127.0.0.1:$PORT/healthz" -o /dev/null 2>/dev/null \
       || curl -sf "http://127.0.0.1:$PORT" -o /dev/null 2>/dev/null; then
        rinfo "code-server is up on 127.0.0.1:$PORT (log: ~/.kiss/code-server.log)."
        exit 0
    fi
    sleep 1
done
echo "ERROR: code-server did not come up on 127.0.0.1:$PORT (see ~/.kiss/code-server.log)" >&2
exit 1
REMOTE
    info "Remote KISS Sorcar / code-server is ready."
fi

# ---------------------------------------------------------------------------
# 4. Open the browser once the tunnel is live, in the background.
# ---------------------------------------------------------------------------
open_browser() {
    for _ in $(seq 1 60); do
        if curl -sf "$URL/healthz" -o /dev/null 2>/dev/null \
           || curl -sf "$URL" -o /dev/null 2>/dev/null; then
            break
        fi
        sleep 1
    done
    info "KISS Sorcar (code-server) is ready at $URL"
    case "$(uname -s)" in
        Darwin) open "$URL" >/dev/null 2>&1 || true ;;
        Linux)  command -v xdg-open &>/dev/null && (xdg-open "$URL" >/dev/null 2>&1 &) || true ;;
    esac
}
open_browser &

# ---------------------------------------------------------------------------
# 5. Hold the SSH tunnel in the foreground.  Ctrl+C closes it (and your
#    connection to the remote code-server); the remote code-server keeps
#    running so you can reconnect by re-running this script.
# ---------------------------------------------------------------------------
echo ""
echo "  ┌──────────────────────────────────────────────────┐"
echo "  │  KISS Sorcar (remote code-server): ${URL}"
echo "  │  Remote host:  ${TARGET}"
echo "  │  Stop tunnel:  Ctrl+C (remote code-server keeps running)"
echo "  │  Remote log:   ~/.kiss/code-server.log on ${TARGET}"
echo "  └──────────────────────────────────────────────────┘"
echo ""
step "Opening SSH tunnel ${URL} -> ${TARGET}:127.0.0.1:${PORT} ..."
exec ssh -N -L "${PORT}:127.0.0.1:${PORT}" "$TARGET"
