#!/bin/bash
# Run on the LOCAL client. Opens an SSH local forward to the vLLM server on the
# Colab VM and runs one term-llm request through it, in the same shell.
#
# Why one shell: in my sandbox, background processes started by one tool call
# are reaped before the next one, so the tunnel and the client must share a
# shell invocation. The trap kills the tunnel on exit either way.
#
# Expects, next to this script:
#   id_colab            dedicated SSH key (ssh-keygen -t ed25519 -f ./id_colab -N '' -C colab-flashnext)
#   term-llm-config/config.yaml   the isolated provider config (term-llm-config.yaml in this bundle)
#
# Environment:
#   COLAB_SESSION   Colab session name (default: flashnext)
#   LOCAL_PORT      local forwarded port (default: 18000)
set -u

SESSION="${COLAB_SESSION:-flashnext}"
LOCAL_PORT="${LOCAL_PORT:-18000}"
HERE="$(cd "$(dirname "$0")" && pwd)"
KEY="$HERE/id_colab"
COLAB_BIN="$(command -v colab)"

export XDG_CONFIG_HOME="$HERE/term-llm-config"
export XDG_DATA_HOME="$HERE/term-llm-data"
mkdir -p "$XDG_DATA_HOME"

# The SSH user is root. 'colab-runtime' is a placeholder host name consumed by
# the ProxyCommand; the Colab CLI bridges SSH over its authenticated WebSocket.
# Host keys change on every runtime, so host-key pinning is disabled here; the
# transport is already authenticated by the Colab CLI's OAuth session.
ssh -i "$KEY" -o IdentitiesOnly=yes \
  -o ProxyCommand="$COLAB_BIN ssh --proxy-mode -s $SESSION --identity $KEY" \
  -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR \
  -o ServerAliveInterval=15 -o ExitOnForwardFailure=yes \
  -N -L "$LOCAL_PORT:127.0.0.1:8000" root@colab-runtime > "$HERE/tunnel.log" 2>&1 &
TUNNEL_PID=$!
trap 'kill $TUNNEL_PID 2>/dev/null; wait $TUNNEL_PID 2>/dev/null' EXIT

for i in $(seq 1 30); do
  if curl -s -m 5 "http://127.0.0.1:$LOCAL_PORT/v1/models" >/dev/null 2>&1; then
    echo "tunnel ready after ${i}s"
    break
  fi
  sleep 1
done

term-llm --stats --no-session ask -p colab --no-search --text --porcelain \
  --max-output-tokens 1024 \
  "Reply with exactly this sentence and nothing else: Hello from Qwen3.8-Flash-Next on Colab G4, served through term-llm."