#!/usr/bin/env bash
# ═══════════════════════════════════════════════════════════════════════════
#  ALFRED AGENT — Upgrade Safeguard Script
#  
#  Handles:
#  1. Pre-upgrade backup of all source, config, data
#  2. Post-upgrade branding re-application (code-server)
#  3. Safe PM2 restart with rollback on failure
#  4. Health check verification
#
#  Usage:
#    ./upgrade.sh backup          # Create timestamped backup
#    ./upgrade.sh brand           # Re-apply code-server branding patches
#    ./upgrade.sh restart         # Safe PM2 restart with health check
#    ./upgrade.sh full            # All three in sequence
#    ./upgrade.sh rollback        # Restore from latest backup
#    ./upgrade.sh status          # Show current agent status
# ═══════════════════════════════════════════════════════════════════════════
set -euo pipefail

AGENT_DIR="$HOME/alfred-agent"
BACKUP_ROOT="$HOME/backups/alfred-agent"
CS_DIR="$HOME/.local/share/code-server/lib/vscode"
HEALTH_URL="http://127.0.0.1:3102/health"
PM2_ID=80
TS=$(date +%Y%m%d-%H%M%S)
BACKUP_DIR="$BACKUP_ROOT/$TS"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log()  { echo -e "${CYAN}[Alfred]${NC} $1"; }
ok()   { echo -e "${GREEN}  ✓${NC} $1"; }
warn() { echo -e "${YELLOW}  ⚠${NC} $1"; }
fail() { echo -e "${RED}  ✗${NC} $1"; }

# ── Backup ────────────────────────────────────────────────────────────
do_backup() {
  log "Creating backup at $BACKUP_DIR"
  mkdir -p "$BACKUP_DIR"
  
  # Agent source
  cp -r "$AGENT_DIR/src" "$BACKUP_DIR/src"
  ok "Source code backed up ($(find "$AGENT_DIR/src" -name '*.js' | wc -l) files)"
  
  # Config files
  cp "$AGENT_DIR/package.json" "$BACKUP_DIR/"
  cp "$AGENT_DIR/ecosystem.config.cjs" "$BACKUP_DIR/"
  [ -f "$AGENT_DIR/.gitignore" ] && cp "$AGENT_DIR/.gitignore" "$BACKUP_DIR/"
  ok "Config files backed up"
  
  # Session data (important — don't lose conversation history)
  if [ -d "$AGENT_DIR/data" ]; then
    cp -r "$AGENT_DIR/data" "$BACKUP_DIR/data"
    local session_count=$(ls "$AGENT_DIR/data/sessions/" 2>/dev/null | wc -l)
    ok "Data directory backed up ($session_count sessions)"
  fi
  
  # Record version info
  cat > "$BACKUP_DIR/MANIFEST.txt" <<EOF
Alfred Agent Backup — $TS
Version: $(node -e "console.log(JSON.parse(require('fs').readFileSync('$AGENT_DIR/package.json','utf8')).version)" 2>/dev/null || echo "unknown")
Node: $(node -v)
Source files: $(find "$AGENT_DIR/src" -name '*.js' | wc -l)
Total lines: $(find "$AGENT_DIR/src" -name '*.js' -exec cat {} + | wc -l)
Services: $(ls "$AGENT_DIR/src/services/" 2>/dev/null | wc -l)
Sessions: $(ls "$AGENT_DIR/data/sessions/" 2>/dev/null | wc -l)
EOF
  ok "Manifest written"
  
  # Prune old backups (keep last 10)
  local count=$(ls -d "$BACKUP_ROOT"/20* 2>/dev/null | wc -l)
  if [ "$count" -gt 10 ]; then
    ls -d "$BACKUP_ROOT"/20* | head -n $(( count - 10 )) | xargs rm -rf
    warn "Pruned old backups (keeping last 10)"
  fi
  
  log "Backup complete: $BACKUP_DIR"
}

# ── Branding ─────────────────────────────────────────────────────────
do_brand() {
  log "Re-applying Alfred IDE branding patches to code-server..."
  
  local WB="$CS_DIR/out/vs/workbench/workbench.web.main.js"
  local NLS_JS="$CS_DIR/out/nls.messages.js"
  local NLS_JSON="$CS_DIR/out/nls.messages.json"
  
  if [ ! -f "$WB" ]; then
    fail "workbench.js not found at $WB — is code-server installed?"
    return 1
  fi
  
  # Backup before patching
  [ ! -f "$WB.bak" ] && cp "$WB" "$WB.bak"
  [ -f "$NLS_JS" ] && [ ! -f "$NLS_JS.bak" ] && cp "$NLS_JS" "$NLS_JS.bak"
  [ -f "$NLS_JSON" ] && [ ! -f "$NLS_JSON.bak" ] && cp "$NLS_JSON" "$NLS_JSON.bak"
  
  # Workbench.js patches (About dialog, menus, notifications)
  sed -i 's/nameShort:"code-server"/nameShort:"Alfred IDE"/g' "$WB"
  sed -i 's/nameLong:"code-server"/nameLong:"Alfred IDE"/g' "$WB"
  # Secure context warning + logout menu
  sed -i 's/d(3228,null,"code-server")/d(3228,null,"Alfred IDE")/g' "$WB"
  sed -i 's/d(3230,null,"code-server")/d(3230,null,"Alfred IDE")/g' "$WB"
  # Update notification
  sed -i 's/\[code-server v/[Alfred IDE v/g' "$WB"
  
  local wb_count=$(grep -c "Alfred IDE" "$WB" 2>/dev/null || echo 0)
  ok "workbench.js: $wb_count 'Alfred IDE' references applied"
  
  # NLS patches (Welcome page, walkthrough headers)
  if [ -f "$NLS_JS" ]; then
    sed -i 's/Get Started with VS Code for the Web/Get Started with Alfred IDE/g' "$NLS_JS"
    sed -i 's/Get Started with VS Code/Get Started with Alfred IDE/g' "$NLS_JS"
    ok "nls.messages.js patched"
  fi
  if [ -f "$NLS_JSON" ]; then
    sed -i 's/Get Started with VS Code for the Web/Get Started with Alfred IDE/g' "$NLS_JSON"
    sed -i 's/Get Started with VS Code/Get Started with Alfred IDE/g' "$NLS_JSON"
    ok "nls.messages.json patched"
  fi
  
  log "Branding patches applied. Restart code-server (PM2 35) to take effect."
}

# ── PM2 Restart with Health Check ────────────────────────────────────
do_restart() {
  log "Restarting Alfred Agent (PM2 $PM2_ID)..."
  
  # Syntax check first
  if ! node -e "import('$AGENT_DIR/src/index.js').then(() => setTimeout(() => process.exit(0), 2000)).catch(e => { console.error(e.message); process.exit(1); })" 2>/dev/null; then
    fail "Import check failed — NOT restarting. Fix errors first."
    return 1
  fi
  ok "Import check passed"
  
  pm2 restart $PM2_ID --update-env 2>/dev/null
  ok "PM2 restart issued"
  
  # Wait and health check
  sleep 3
  local health
  health=$(curl -s --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "")
  
  if echo "$health" | grep -q '"status":"online"'; then
    local version=$(echo "$health" | python3 -c "import sys,json; print(json.load(sys.stdin).get('version','?'))" 2>/dev/null || echo "?")
    ok "Health check passed — v$version online"
  else
    fail "Health check FAILED — agent may not be running"
    warn "Check logs: pm2 logs $PM2_ID --lines 20"
    return 1
  fi
}

# ── Rollback ─────────────────────────────────────────────────────────
do_rollback() {
  local latest=$(ls -d "$BACKUP_ROOT"/20* 2>/dev/null | tail -1)
  if [ -z "$latest" ]; then
    fail "No backups found in $BACKUP_ROOT"
    return 1
  fi
  
  log "Rolling back from $latest"
  
  # Restore source
  rm -rf "$AGENT_DIR/src"
  cp -r "$latest/src" "$AGENT_DIR/src"
  ok "Source restored"
  
  # Restore config
  cp "$latest/package.json" "$AGENT_DIR/"
  cp "$latest/ecosystem.config.cjs" "$AGENT_DIR/"
  ok "Config restored"
  
  do_restart
  log "Rollback complete"
}

# ── Status ───────────────────────────────────────────────────────────
do_status() {
  echo ""
  echo -e "${CYAN}╔═══════════════════════════════════════════════════════════╗${NC}"
  echo -e "${CYAN}║  ALFRED AGENT — System Status                           ║${NC}"
  echo -e "${CYAN}╚═══════════════════════════════════════════════════════════╝${NC}"
  echo ""
  
  # Health
  local health
  health=$(curl -s --max-time 5 "$HEALTH_URL" 2>/dev/null || echo '{"status":"offline"}')
  local status=$(echo "$health" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('status','offline'))" 2>/dev/null || echo "offline")
  local version=$(echo "$health" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('version','?'))" 2>/dev/null || echo "?")
  local sessions=$(echo "$health" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('activeSessions',0))" 2>/dev/null || echo "0")
  local features=$(echo "$health" | python3 -c "import sys,json; d=json.load(sys.stdin); print(', '.join(d.get('features',[])))" 2>/dev/null || echo "?")
  
  if [ "$status" = "online" ]; then
    echo -e "  Status:   ${GREEN}● ONLINE${NC}"
  else
    echo -e "  Status:   ${RED}● OFFLINE${NC}"
  fi
  echo "  Version:  $version"
  echo "  Sessions: $sessions active"
  echo "  Features: $features"
  echo ""
  
  # Source stats
  local src_files=$(find "$AGENT_DIR/src" -name '*.js' | wc -l)
  local src_lines=$(find "$AGENT_DIR/src" -name '*.js' -exec cat {} + | wc -l)
  local svc_count=$(ls "$AGENT_DIR/src/services/" 2>/dev/null | wc -l)
  echo "  Source:   $src_files files, $src_lines lines"
  echo "  Services: $svc_count modules"
  echo ""
  
  # GoForge Git
  if [ -d "$AGENT_DIR/.git" ]; then
    local branch=$(cd "$AGENT_DIR" && git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "?")
    local dirty=$(cd "$AGENT_DIR" && git status --porcelain 2>/dev/null | wc -l)
    local last_commit=$(cd "$AGENT_DIR" && git log --oneline -1 2>/dev/null || echo "no commits")
    echo "  GoForge:  branch=$branch, $dirty dirty files"
    echo "  Last:     $last_commit"
    echo ""
  fi
  
  # Backups
  local backup_count=$(ls -d "$BACKUP_ROOT"/20* 2>/dev/null | wc -l)
  local latest_backup=$(ls -d "$BACKUP_ROOT"/20* 2>/dev/null | tail -1 | xargs basename 2>/dev/null || echo "none")
  echo "  Backups:  $backup_count saved (latest: $latest_backup)"
  echo ""
  
  # PM2
  echo "  PM2 ID:   $PM2_ID"
  pm2 show $PM2_ID 2>/dev/null | grep -E "status|uptime|restart" | head -3 | sed 's/^/  /'
  echo ""
}

# ── Full Upgrade ─────────────────────────────────────────────────────
do_full() {
  log "Running full upgrade sequence..."
  echo ""
  do_backup
  echo ""
  do_brand
  echo ""
  do_restart
  echo ""
  log "Full upgrade sequence complete."
}

# ── Main ─────────────────────────────────────────────────────────────
case "${1:-help}" in
  backup)   do_backup ;;
  brand)    do_brand ;;
  restart)  do_restart ;;
  full)     do_full ;;
  rollback) do_rollback ;;
  status)   do_status ;;
  *)
    echo "Alfred Agent Upgrade Safeguard"
    echo ""
    echo "Usage: $0 {backup|brand|restart|full|rollback|status}"
    echo ""
    echo "  backup    Create timestamped backup of source, config, data"
    echo "  brand     Re-apply Alfred IDE branding to code-server"
    echo "  restart   Safe PM2 restart with health check"
    echo "  full      backup + brand + restart (for after code-server upgrades)"
    echo "  rollback  Restore from latest backup"
    echo "  status    Show agent status, features, GoForge, backups"
    ;;
esac