#!/usr/bin/env bash
# bootstrap.sh — one-shot setup for the Agent0 box.
#
# Run once after cloning. Handles everything that can be automated:
#   - Ollama install + model pulls
#   - SSH key generation (tunnel key + Gitea key)
#   - gutasktool clone (into agent-zero-data, mounted into the agent0 container)
#   - Docker containers (gutasktool installed inside agent0 container)
#
# After this script: add two SSH keys (printed at the end), then open
# https://agent0.glitch.university and enter your API keys in Settings.

set -euo pipefail

REPO_DIR="$(cd "$(dirname "$0")" && pwd)"
OS="$(uname -s)"
GUTASK_DIR="${REPO_DIR}/agent-zero-data/gutasktool"
GITEA_HOST="ramanujan.glitch.university"
GITEA_PORT=2222
VPS_HOST="glitch.university"
VPS_TUNNEL_USER="tunnel"
TUNNEL_KEY="${REPO_DIR}/tunnel/id_ed25519"
GITEA_KEY="${HOME}/.ssh/gitea_ed25519"

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; CYAN='\033[0;36m'; NC='\033[0m'
info()    { echo -e "${GREEN}✓${NC} $*"; }
prompt()  { echo -e "${CYAN}?${NC} $*"; }
section() { echo ""; echo -e "${YELLOW}── $* ──────────────────────────────────────────${NC}"; }

# ── Prerequisites ─────────────────────────────────────────────────────────────
section "Checking prerequisites"

command -v docker      >/dev/null 2>&1 || { echo "docker not found. Install Docker first."; exit 1; }
command -v ssh-keygen  >/dev/null 2>&1 || { echo "ssh-keygen not found."; exit 1; }
command -v ssh-keyscan >/dev/null 2>&1 || { echo "ssh-keyscan not found."; exit 1; }
docker compose version >/dev/null 2>&1 || { echo "docker compose not found."; exit 1; }

info "Prerequisites OK"

# ── Credentials (collected upfront before any private repo clones) ─────────────
section "Credentials"

echo ""
prompt "Enter CONTENT_API_KEY for glitch.university (find it in the VPS .env): "
read -r -s CONTENT_API_KEY
echo ""

echo ""
prompt "Enter Gitea personal access token (needs repo read access — create at https://${GITEA_HOST}/user/settings/applications): "
read -r -s GITEA_TOKEN_BOOTSTRAP
echo ""

info "Credentials collected"

# ── Ollama ────────────────────────────────────────────────────────────────────
section "Ollama"

if command -v ollama >/dev/null 2>&1; then
  info "Ollama already installed"
else
  info "Installing Ollama..."
  curl -fsSL https://ollama.com/install.sh | sh
fi

# Ensure the Ollama server is running
if [[ "$OS" == "Linux" ]]; then
  if systemctl is-active --quiet ollama 2>/dev/null; then
    info "Ollama service running"
  else
    info "Starting Ollama service..."
    sudo systemctl enable --now ollama
  fi
else
  # macOS — check if server is already answering, otherwise start in background
  if curl -sf http://localhost:11434/api/tags >/dev/null 2>&1; then
    info "Ollama server already running"
  else
    info "Starting Ollama server..."
    ollama serve >/tmp/ollama.log 2>&1 &
    sleep 2
    curl -sf http://localhost:11434/api/tags >/dev/null 2>&1 \
      && info "Ollama server started" \
      || echo "  Warning: Ollama server didn't respond — start it manually with: ollama serve"
  fi
fi

# Returns 0 if model is already pulled, 1 if not
model_installed() {
  ollama list 2>/dev/null | awk 'NR>1 {print $1}' | grep -qx "$1"
}

pull_model() {
  local model="$1"
  if model_installed "$model"; then
    info "Already downloaded: $model"
  else
    info "Pulling $model..."
    ollama pull "$model"
  fi
}

echo ""
echo "  ── Coding ──────────────────────────────────────"
echo "  [1]  qwen2.5-coder:7b       ~5 GB   fast, good for quick tasks"
echo "  [2]  qwen2.5-coder:32b      ~20 GB  strong coding (recommended)"
echo "  [3]  qwen2.5-coder:72b      ~45 GB  best coding quality"
echo "  [4]  deepseek-coder-v2:16b  ~9 GB   strong alternative"
echo ""
echo "  ── General ─────────────────────────────────────"
echo "  [5]  llama3.1:8b            ~5 GB   fast general-purpose"
echo "  [6]  llama3.1:70b           ~40 GB  strong general-purpose"
echo "  [7]  mistral:7b             ~4 GB   fast, widely used"
echo "  [8]  mixtral:8x7b           ~26 GB  MoE, strong at instruction following"
echo "  [9]  gemma2:9b              ~6 GB   Google, good at reasoning"
echo "  [10] gemma2:27b             ~16 GB  Google, stronger"
echo "  [11] qwen2.5:72b            ~45 GB  strong general + multilingual"
echo ""
echo "  ── Reasoning ───────────────────────────────────"
echo "  [12] deepseek-r1:7b         ~5 GB   chain-of-thought, fast"
echo "  [13] deepseek-r1:32b        ~20 GB  chain-of-thought, strong"
echo "  [14] deepseek-r1:70b        ~43 GB  chain-of-thought, best"
echo ""
prompt "Which models to pull? (e.g. 1 3 7  or 'all'  or Enter to skip): "
read -r MODEL_CHOICES

# Expand 'all' to full list
if [[ "$MODEL_CHOICES" == "all" ]]; then
  MODEL_CHOICES="1 2 3 4 5 6 7 8 9 10 11 12 13 14"
fi

for choice in $MODEL_CHOICES; do
  case "$choice" in
    1)  pull_model "qwen2.5-coder:7b"       ;;
    2)  pull_model "qwen2.5-coder:32b"      ;;
    3)  pull_model "qwen2.5-coder:72b"      ;;
    4)  pull_model "deepseek-coder-v2:16b"  ;;
    5)  pull_model "llama3.1:8b"            ;;
    6)  pull_model "llama3.1:70b"           ;;
    7)  pull_model "mistral:7b"             ;;
    8)  pull_model "mixtral:8x7b"           ;;
    9)  pull_model "gemma2:9b"              ;;
    10) pull_model "gemma2:27b"             ;;
    11) pull_model "qwen2.5:72b"            ;;
    12) pull_model "deepseek-r1:7b"         ;;
    13) pull_model "deepseek-r1:32b"        ;;
    14) pull_model "deepseek-r1:70b"        ;;
  esac
done

# ── SSH keys ──────────────────────────────────────────────────────────────────
section "SSH key: glitch-tunnel → VPS"

mkdir -p "${REPO_DIR}/tunnel"
chmod 700 "${REPO_DIR}/tunnel"

if [[ -f "$TUNNEL_KEY" ]]; then
  info "Tunnel key already exists"
else
  ssh-keygen -t ed25519 -f "$TUNNEL_KEY" -C "glitch-tunnel@$(hostname)" -N ""
  info "Tunnel key generated at ./tunnel/id_ed25519"
fi

info "Scanning VPS host key..."
ssh-keyscan -p 22 "$VPS_HOST" > "${REPO_DIR}/tunnel/known_hosts" 2>/dev/null \
  || { echo "Could not reach ${VPS_HOST} — check network. Re-run when online."; exit 1; }

section "SSH key: Gitea (git operations)"

mkdir -p "${HOME}/.ssh"
chmod 700 "${HOME}/.ssh"

if [[ -f "$GITEA_KEY" ]]; then
  info "Gitea key already exists (local only — verify it is registered on the server)"
else
  ssh-keygen -t ed25519 -f "$GITEA_KEY" -C "gunnar@$(hostname)" -N ""
  info "Gitea key generated at ${GITEA_KEY}"
fi

echo ""
echo "  This public key must be registered at https://${GITEA_HOST} → Settings → SSH Keys:"
echo ""
echo "  $(cat "${GITEA_KEY}.pub")"
echo ""

# Add Gitea to known_hosts to avoid interactive prompt during git operations
if ! ssh-keygen -F "[${GITEA_HOST}]:${GITEA_PORT}" -f "${HOME}/.ssh/known_hosts" >/dev/null 2>&1; then
  ssh-keyscan -p "$GITEA_PORT" "$GITEA_HOST" >> "${HOME}/.ssh/known_hosts" 2>/dev/null || true
fi

# Write SSH config entry for Gitea
if ! grep -q "Host ${GITEA_HOST}" "${HOME}/.ssh/config" 2>/dev/null; then
  cat >> "${HOME}/.ssh/config" << EOF

Host ${GITEA_HOST}
  HostName ${GITEA_HOST}
  Port ${GITEA_PORT}
  User git
  IdentityFile ${GITEA_KEY}
EOF
  info "SSH config entry added for ${GITEA_HOST}"
fi

# ── gutasktool ────────────────────────────────────────────────────────────────
section "gutasktool"

mkdir -p "${REPO_DIR}/agent-zero-data"

if [[ -d "$GUTASK_DIR" ]]; then
  info "gutasktool already at ${GUTASK_DIR}"
else
  info "Cloning gutasktool..."
  git clone "ssh://git@${GITEA_HOST}:${GITEA_PORT}/glitch-university/gutasktool.git" "$GUTASK_DIR" \
    || git clone "https://oauth2:${GITEA_TOKEN_BOOTSTRAP}@${GITEA_HOST}/glitch-university/gutasktool.git" "$GUTASK_DIR"
fi

if [[ ! -f "${GUTASK_DIR}/.env" ]]; then
  cat > "${GUTASK_DIR}/.env" << EOF
API_URL=https://glitch.university
CONTENT_API_KEY=${CONTENT_API_KEY}
GITEA_URL=https://${GITEA_HOST}
GITEA_OWNER=glitch-university
# AGENT_ID, AGENT_NAME, AGENT_PASSWORD, and GITEA_TOKEN are per-agent.
# They are passed through each agent's thread context, not set here.
EOF
  info "gutasktool .env created at ${GUTASK_DIR}/.env"
else
  info "gutasktool .env already exists"
fi

# ── Docker containers ─────────────────────────────────────────────────────────
section "Docker containers"

cd "$REPO_DIR"
info "Building glitch-tunnel image..."
docker compose build --quiet

info "Starting containers..."
docker compose up -d

info "Installing gutasktool inside agent0 container..."
docker exec agent0 bash -c "apt-get update -q && apt-get install -y -q python3-pip && python3 -m pip install --break-system-packages -e /a0/usr/gutasktool --quiet"

echo ""
docker compose ps

# ── Manual steps summary ──────────────────────────────────────────────────────
TUNNEL_PUBKEY=$(cat "${TUNNEL_KEY}.pub")
GITEA_PUBKEY=$(cat "${GITEA_KEY}.pub")

echo ""
echo -e "${YELLOW}════════════════════════════════════════════════════════${NC}"
echo -e "${YELLOW}  Two manual steps remaining:${NC}"
echo -e "${YELLOW}════════════════════════════════════════════════════════${NC}"
echo ""
echo -e "${CYAN}1. Add tunnel key to VPS${NC}"
echo "   Run on the VPS:"
echo ""
echo "   echo '${TUNNEL_PUBKEY}' >> /home/tunnel/.ssh/authorized_keys"
echo ""
echo -e "${CYAN}2. Add Gitea key to Gitea${NC}"
echo "   Log into https://${GITEA_HOST} as gunnar"
echo "   Settings → SSH Keys → Add key:"
echo ""
echo "   ${GITEA_PUBKEY}"
echo ""
echo -e "${YELLOW}════════════════════════════════════════════════════════${NC}"
echo ""
echo "  After adding both keys:"
echo "  • Tunnel connects automatically (check: docker logs glitch-tunnel)"
echo "  • Open https://agent0.glitch.university"
echo "  • Enter basic auth password"
echo "  • Settings → add Anthropic / OpenAI API keys"
echo "  • Ollama is pre-configured at http://host.docker.internal:11434"
echo ""
echo -e "${GREEN}Bootstrap complete.${NC}"