#!/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 user@ip-address [PORT]
#
# Flow:
#   1. SSH into user@ip-address.
#   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
# ---------------------------------------------------------------------------
TARGET="${1:-}"
PORT="${2:-8080}"

if [[ -z "$TARGET" ]]; then
    die "Usage: $0 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"

# ---------------------------------------------------------------------------
# 3. Remote bootstrap: code-server + kiss_ai + launch code-server.
#
# The heredoc is single-quoted so nothing is expanded locally; PORT is passed
# through the environment on the ssh command line instead.
# ---------------------------------------------------------------------------
step "Provisioning KISS Sorcar on $TARGET (this can take several minutes)..."
ssh "$TARGET" "PORT='$PORT' bash -s" <<'REMOTE'
set -e
PORT="${PORT:-8080}"
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)" "$*"; }

# --- 3a. code-server -------------------------------------------------------
if ! command -v code-server &>/dev/null; then
    rinfo "Installing code-server..."
    curl -fsSL https://code-server.dev/install.sh | sh
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; }
WORKSPACE="$HOME/kiss_ai"

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

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

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